pokouqiege/models/line_detect/cs_readv.py

import os
import csv
import json
import math
import ast
import shutil

def ellipse_to_box_points(cx, cy, rx, ry, theta):
    """
    ½«ÍÖÔ²²ÎÊýת»»Îª¾ØÐΣ¨ÉÏÏÂ×óÓÒ4µã£©
    ½üËÆ·½·¨£ºÓÃÐýת¾ØÐεÄ4¸ö¶¥µã±íʾÍÖÔ²±ß½ç
    """
    # Ëĸöµã£¨×óÓÒÖе㣬ÉÏÏÂÖе㣩
    points = [
        [cx - rx, cy],  # ×ó
        [cx + rx, cy],  # ÓÒ
        [cx, cy - ry],  # ÉÏ
        [cx, cy + ry],  # ÏÂ
    ]

    # ÐýתÕâЩµã
    cos_t = math.cos(theta)
    sin_t = math.sin(theta)
    rotated_points = []
    for x, y in points:
        x_new = (x - cx) * cos_t - (y - cy) * sin_t + cx
        y_new = (x - cx) * sin_t + (y - cy) * cos_t + cy
        rotated_points.append([x_new, y_new])

    return rotated_points


def csv_to_labelme_json(csv_path, img_folder, output_folder, image_exts=(".png", ".tiff")):
    os.makedirs(output_folder, exist_ok=True)

    with open(csv_path, newline='', encoding='utf-8') as csvfile:
        reader = csv.DictReader(csvfile)
        for row in reader:
            filename = row["filename"]
            region_shape_attributes = row["region_shape_attributes"]

            # ת³É×ֵ䣨CSV ÀïÊÇ×Ö·û´®£©
            try:
                attr = json.loads(region_shape_attributes)
            except json.JSONDecodeError:
                attr = ast.literal_eval(region_shape_attributes)

            if attr.get("name") != "ellipse":
                continue

            cx, cy = float(attr["cx"]), float(attr["cy"])
            rx, ry = float(attr["rx"]), float(attr["ry"])
            theta = float(attr["theta"])

            # ÍÖÔ²Ëĸöµã
            points = ellipse_to_box_points(cx, cy, rx, ry, theta)

            # »ñÈ¡xmin, ymin, xmax, ymax
            xs = [p[0] for p in points]
            ys = [p[1] for p in points]
            xmin, ymin, xmax, ymax = int(min(xs)), int(min(ys)), int(max(xs)), int(max(ys))

            # Éú³É Labelme ¸ñʽ
            data = {
                "version": "5.0.1",
                "flags": {},
                "shapes": [
                    {
                        "label": "circle",
                        "points": points,
                        "shape_type": "polygon",
                        "flags": {},
                        "xmin": xmin,
                        "ymin": ymin,
                        "xmax": xmax,
                        "ymax": ymax,
                    }
                ],
                "imagePath": filename.replace(".png", ".jpg").replace(".tiff", ".jpg"),
                "imageHeight": 2000,  # Èç¹ûÄãÖªµÀʵ¼ÊͼÏñ³ß´ç¿ÉÒÔ¸Ä
                "imageWidth": 2000
            }

            # дÈë JSON Îļþ
            json_filename = os.path.splitext(filename)[0] + ".json"
            json_path = os.path.join(output_folder, json_filename)
            with open(json_path, "w", encoding="utf-8") as f:
                json.dump(data, f, indent=4, ensure_ascii=False)

            # ¿½±´Í¼Æ¬Îļþ
            src_img_path = None
            for ext in image_exts:
                test_path = os.path.join(img_folder, filename.replace(".png", ext).replace(".tiff", ext))
                if os.path.exists(test_path):
                    src_img_path = test_path
                    break

            if src_img_path:
                dst_img_path = os.path.join(output_folder, os.path.basename(src_img_path))
                shutil.copy(src_img_path, dst_img_path)
            else:
                print(f"?? ÕÒ²»µ½Í¼Æ¬: {filename}")

    print(f"? ת»»Íê³É£¬Êä³ö·¾¶£º{output_folder}")


if __name__ == "__main__":
    # ÊäÈëÎļþ¼Ð
    csv_folder = "/data/share/zyh/数据集汇总/circle/huayan_circle/csv"
    img_folder = "/data/share/zyh/数据集汇总/circle/huayan_circle/彩色图"
    output_folder = "/home/zhaoyinghan/py_ws/data/circle/huayan"

    for csv_file in os.listdir(csv_folder):
        if csv_file.endswith(".csv"):
            csv_path = os.path.join(csv_folder, csv_file)
            csv_to_labelme_json(csv_path, img_folder, output_folder)