import json
from pathlib import Path
from PIL import Image

# ======================
# Paths Configuration
# ======================
JSON_DIR = Path("/data/share/zyh/master_dataset/circle/huayan_circle/251112_251111/day02_machine")
IMAGE_DIR = Path("/data/share/zyh/master_dataset/circle/huayan_circle/251112_251111/day02_machine")  # source images
OUTPUT_DIR = JSON_DIR  # overwrite original or use a new folder

# ======================
# Process each JSON
# ======================
for json_file in JSON_DIR.glob("*.json"):
    with open(json_file, "r") as f:
        data = json.load(f)

    # Determine corresponding image file
    stem_name = json_file.stem
    if stem_name.endswith("_annotations"):
        stem_name = stem_name.replace("_annotations", "")
    possible_extensions = [".png", ".jpg", ".jpeg"]
    img_file = None
    for ext in possible_extensions:
        candidate = IMAGE_DIR / f"{stem_name}{ext}"
        if candidate.exists():
            img_file = candidate
            break
    if img_file is None:
        print(f"[WARNING] No image found for JSON: {json_file.name}")
        continue

    # Load image size
    image = Image.open(img_file)
    img_width, img_height = image.size

    # Convert each entry to LabelMe shape (without arc)
    shapes = []
    for ann in data if isinstance(data, list) else data.get("annotations", []):
        shape = {
            "label": ann.get("label", ann.get("class_name", "circle")),
            "points": ann.get("points", []),
            "shape_type": "polygon",
            "flags": {},
            "xmin": ann.get("xmin", 0),
            "ymin": ann.get("ymin", 0),
            "xmax": ann.get("xmax", 0),
            "ymax": ann.get("ymax", 0)
        }
        shapes.append(shape)

    # Create LabelMe JSON
    labelme_json = {
        "version": "5.0.1",
        "flags": {},
        "shapes": shapes,
        "imagePath": img_file.name,
        "imageHeight": img_height,
        "imageWidth": img_width
    }

    # Save updated JSON (overwrite or to another folder)
    output_path = OUTPUT_DIR / json_file.name
    with open(output_path, "w") as f:
        json.dump(labelme_json, f, indent=4)

    print(f"Processed JSON saved: {output_path.name}")

print("\n? All JSON files converted to LabelMe format (arc removed).")