pokouqiege/utils/data_process/zyh/tuoyuan_rgb/a_longshort.py

import json
import os
import cv2
import shutil

# 输入文件夹
json_folder = r"\\192.168.50.222\share\zyh\data\via\via_json"
image_folder = r"\\192.168.50.222\share\zqy\colorphoto_for_guanban"
output_folder = r"\\192.168.50.222\share\zyh\data\rgb_tuoyuan_axis_bbox\source"

os.makedirs(output_folder, exist_ok=True)

for json_file in os.listdir(json_folder):
    if not json_file.lower().endswith(".json"):
        continue

    json_path = os.path.join(json_folder, json_file)

    with open(json_path, "r", encoding="utf-8") as f:
        data = json.load(f)

    img_metadata = data.get("_via_img_metadata", {})

    for img_info in img_metadata.values():
        filename = img_info.get("filename")
        regions = img_info.get("regions", [])

        # 只取椭圆
        ellipses = [r for r in regions if r.get("shape_attributes", {}).get("name") == "ellipse"]
        if not ellipses:
            continue

        # 找图片路径
        img_path = os.path.join(image_folder, filename)
        if not os.path.exists(img_path):
            print(f"图片未找到: {filename}")
            continue

        # 读取图片尺寸
        img = cv2.imread(img_path)
        if img is None:
            print(f"图片无法读取: {filename}")
            continue
        height, width = img.shape[:2]

        # 构造输出 JSON
        output_data = {
            "version": "5.0.1",
            "flags": {},
            "shapes": [],
            "imagePath": filename,
            "imageHeight": height,
            "imageWidth": width
        }

        for e in ellipses:
            cx = e["shape_attributes"]["cx"]
            cy = e["shape_attributes"]["cy"]
            rx = e["shape_attributes"]["rx"]
            ry = e["shape_attributes"]["ry"]

            # 四个端点
            left = [cx - rx, cy]
            right = [cx + rx, cy]
            top = [cx, cy - ry]
            bottom = [cx, cy + ry]

            # bbox
            xmin = min(left[0], right[0], top[0], bottom[0])
            xmax = max(left[0], right[0], top[0], bottom[0])
            ymin = min(left[1], right[1], top[1], bottom[1])
            ymax = max(left[1], right[1], top[1], bottom[1])

            # 长轴（水平）
            long_axis = [left, right]
            # 短轴（垂直）
            short_axis = [top, bottom]

            # 保存长轴
            output_data["shapes"].append({
                "label": "long_axis",
                "points": long_axis,
                "shape_type": "line",
                "flags": {},
                "xmin": int(xmin),
                "ymin": int(ymin),
                "xmax": int(xmax),
                "ymax": int(ymax)
            })

            # 保存短轴
            output_data["shapes"].append({
                "label": "short_axis",
                "points": short_axis,
                "shape_type": "line",
                "flags": {},
                "xmin": int(xmin),
                "ymin": int(ymin),
                "xmax": int(xmax),
                "ymax": int(ymax)
            })

        # 保存 JSON
        json_output_path = os.path.join(output_folder, filename.replace(".jpg", ".json"))
        with open(json_output_path, "w", encoding="utf-8") as out_f:
            json.dump(output_data, out_f, ensure_ascii=False, indent=4)

        # 复制图片
        img_output_path = os.path.join(output_folder, filename)
        shutil.copy(img_path, img_output_path)

        print(f"已生成: {json_output_path} 和 {img_output_path}")