pokouqiege/utils/data_process/zyh/tiff_change_pic_reshape.py

import os
import json
from itertools import combinations
import numpy as np
import skimage
from scipy.ndimage import zoom
from tqdm import tqdm
import imageio.v3 as iio

def to_int(x):
    return tuple(map(int, x))

def convert_labelme_to_custom_json(labelme_json_path):
    try:
        with open(labelme_json_path, 'r', encoding='utf-8') as f:
            labelme_data = json.load(f)
        if 'imageHeight' not in labelme_data:
            raise KeyError("Missing 'imageHeight' in JSON")
        if 'imageWidth' not in labelme_data:
            raise KeyError("Missing 'imageWidth' in JSON")
        image_height = labelme_data['imageHeight']
        image_width = labelme_data['imageWidth']
        lines = []
        for shape in labelme_data.get('shapes', []):
            if shape.get('label') == 'dseam1':
                if len(shape.get('points', [])) != 2:
                    print(f"Warning: Skipping invalid line in {labelme_json_path}")
                    continue
                start_point = shape['points'][0]
                end_point = shape['points'][1]
                lines.append([start_point, end_point])
        if not lines:
            raise ValueError("No line segments found in the JSON")
        custom_json_data = {
            "image_id": os.path.splitext(os.path.basename(labelme_json_path))[0],
            "lines": lines
        }
        return custom_json_data
    except Exception as e:
        print(f"Error reading {labelme_json_path}: {e}")
        return None

def find_image_file(json_path, image_extensions=('.jpg', '.jpeg', '.png', '.bmp', '.tiff')):
    base_name = os.path.splitext(json_path)[0]
    for ext in image_extensions:
        image_path = f"{base_name}{ext}"
        if os.path.exists(image_path):
            return image_path
    return None

def generate_heatmap_data(image, lines):
    try:
        im_rescale = (512, 512)
        heatmap_scale = (128, 128)
        fy, fx = heatmap_scale[1] / image.shape[0], heatmap_scale[0] / image.shape[1]
        jmap = np.zeros((1,) + heatmap_scale, dtype=np.float32)
        joff = np.zeros((1, 2) + heatmap_scale, dtype=np.float32)
        lmap = np.zeros(heatmap_scale, dtype=np.float32)
        lines[:, :, 0] = np.clip(lines[:, :, 0] * fx, 0, heatmap_scale[0] - 1e-4)
        lines[:, :, 1] = np.clip(lines[:, :, 1] * fy, 0, heatmap_scale[1] - 1e-4)
        lines = lines[:, :, ::-1]
        junc = []
        jids = {}
        def jid(jun):
            jun = tuple(jun[:2])
            if jun in jids:
                return jids[jun]
            jids[jun] = len(junc)
            junc.append(np.array(jun + (0,)))
            return len(junc) - 1
        lnid = []
        lpos, lneg = [], []
        for v0, v1 in lines:
            lnid.append((jid(v0), jid(v1)))
            lpos.append([junc[jid(v0)], junc[jid(v1)]])
            vint0, vint1 = to_int(v0), to_int(v1)
            jmap[0][vint0] = 1
            jmap[0][vint1] = 1
            rr, cc, value = skimage.draw.line_aa(*to_int(v0), *to_int(v1))
            lmap[rr, cc] = np.maximum(lmap[rr, cc], value)
        for v in junc:
            vint = to_int(v[:2])
            joff[0, :, vint[0], vint[1]] = v[:2] - vint - 0.5
        llmap = zoom(lmap, [0.5, 0.5])
        lineset = set([frozenset(l) for l in lnid])
        for i0, i1 in combinations(range(len(junc)), 2):
            if frozenset([i0, i1]) not in lineset:
                v0, v1 = junc[i0], junc[i1]
                vint0, vint1 = to_int(v0[:2] / 2), to_int(v1[:2] / 2)
                rr, cc, value = skimage.draw.line_aa(*vint0, *vint1)
                lneg.append([v0, v1, i0, i1, np.average(np.minimum(value, llmap[rr, cc]))])
        if len(lneg) == 0:
            for i0, i1 in combinations(range(len(junc)), 2):
                v0, v1 = junc[i0], junc[i1]
                vint0, vint1 = to_int(v0[:2] / 2), to_int(v1[:2] / 2)
                rr, cc, value = skimage.draw.line_aa(*vint0, *vint1)
                simulated_value = np.random.uniform(0.01, 0.1)
                lneg.append([v0, v1, i0, i1, simulated_value])
        lneg.sort(key=lambda l: -l[-1])
        junc = np.array(junc, dtype=np.float32)
        Lpos = np.array(lnid, dtype=np.int32)
        Lneg = np.array([l[2:4] for l in lneg][:4000], dtype=np.int32)
        lpos = np.array(lpos, dtype=np.float32)
        lneg = np.array([l[:2] for l in lneg[:2000]], dtype=np.float32)
        return {
            "junc_map": {"name": "junc_map", "shape": list(jmap.shape), "content": jmap.tolist()},
            "junc_offset": {"name": "junc_offset", "shape": list(joff.shape), "content": joff.tolist()},
            "line_map": {"name": "line_map", "shape": list(lmap.shape), "content": lmap.tolist()},
            "junc_coords": {"name": "junc_coords", "shape": list(junc.shape), "content": junc.tolist()},
            "line_pos_idx": {"name": "line_pos_idx", "shape": list(Lpos.shape), "content": Lpos.tolist()},
            "line_neg_idx": {"name": "line_neg_idx", "shape": list(Lneg.shape), "content": Lneg.tolist()},
            "line_pos_coords": {"name": "line_pos_coords", "shape": list(lpos.shape), "content": lpos.tolist()},
            "line_neg_coords": {"name": "line_neg_coords", "shape": list(lneg.shape), "content": lneg.tolist()}
        }
    except Exception as e:
        print(f"Error generating heatmap data: {e}")
        raise

def batch_process(labelme_json_dir, output_dir):
    json_files = [f for f in os.listdir(labelme_json_dir) if f.endswith('.json')]
    os.makedirs(output_dir, exist_ok=True)
    for filename in tqdm(json_files, desc="Processing JSON files"):
        labelme_json_path = os.path.join(labelme_json_dir, filename)
        try:
            custom_json_data = convert_labelme_to_custom_json(labelme_json_path)
            if custom_json_data is None:
                continue
            image_path = find_image_file(labelme_json_path)
            if not image_path or not os.path.exists(image_path):
                print(f"Image not found: {labelme_json_path}")
                continue
            image = iio.imread(image_path)
            if image.ndim == 2:
                image = image[:, :, np.newaxis]
            elif image.ndim == 3 and image.shape[2] == 4:
                image = image[:, :, :3]  # Drop alpha if needed

            # 保存缩放后 TIFF（保持原始位深度）
            zoom_factors = [512 / image.shape[0], 512 / image.shape[1], 1]

            resized_image = zoom(image, zoom_factors, order=0)
            
            resized_image = resized_image.astype(image.dtype)
            image_save_path = os.path.join(output_dir, os.path.splitext(filename)[0] + '.tiff')
            iio.imwrite(image_save_path, resized_image)

            # heatmap 处理
            lines = np.array(custom_json_data['lines']).reshape(-1, 2, 2)
            wires_data = generate_heatmap_data(image, lines)

            # 输出 JSON 文件
            custom_json_data.pop("lines", None)
            custom_json_data.update({
                "boxes": [],
                "segmentations": [],
                "wires": [wires_data]
            })
            output_json_path = os.path.join(output_dir, os.path.splitext(filename)[0] + '.json')
            with open(output_json_path, 'w', encoding='utf-8') as f:
                json.dump(custom_json_data, f, separators=(',', ':'))
        except Exception as e:
            print(f"Error processing {labelme_json_path}: {e}")

if __name__ == "__main__":
    labelme_json_dir = r"\\192.168.50.222\share\zyh\512\init\paomo"
    custom_json_output_dir = r"\\192.168.50.222\share\zyh\512\init\target"
    batch_process(labelme_json_dir, custom_json_output_dir)