# Run Batch_Inference.py to generate image–mask pairs. 
# Use this script to convert the image–mask pairs into a COCO 1.1-compatible JSON annotation file. 
# The resulting JSON can be imported into CVAT as annotations.

import os
import cv2
import json
from tqdm import tqdm
from PIL import Image

# Input folders
images_dir = "newimgs/images"   
masks_dir = "newimgs/masks"
output_json = "annotations.json"

# COCO format template
coco = {
    "images": [],
    "annotations": [],
    "categories": [{"id": 1, "name": "leaf"}]  # change "leaf" if needed
}

ann_id = 1
img_id = 1

for fname in tqdm(os.listdir(images_dir)):
    if fname.lower().endswith((".jpg", ".jpeg", ".png")):
        name = os.path.splitext(fname)[0]
        img_path = os.path.join(images_dir, fname)
        mask_path = os.path.join(masks_dir, name + ".png")

        if not os.path.exists(mask_path):
            print(f" No mask for {fname}, skipping...")
            continue

        # Load image and mask
        img = Image.open(img_path).convert("RGB")
        w, h = img.size
        mask = cv2.imread(mask_path, 0)

        # Ensure binary
        _, mask_bin = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)

        # Add image info
        coco["images"].append({
            "id": img_id,
            "file_name": fname,
            "width": w,
            "height": h
        })

        # Extract polygons from mask
        contours, _ = cv2.findContours(mask_bin, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        for c in contours:
            if len(c) < 6:  # too few points
                continue

            segmentation = c.flatten().tolist()
            x, y, w_box, h_box = cv2.boundingRect(c)

            coco["annotations"].append({
                "id": ann_id,
                "image_id": img_id,
                "category_id": 1,
                "segmentation": [segmentation],
                "bbox": [x, y, w_box, h_box],
                "area": float(cv2.contourArea(c)),
                "iscrowd": 0
            })
            ann_id += 1

        img_id += 1

# Save JSON
with open(output_json, "w") as f:
    json.dump(coco, f)

print(f"Done. Saved {len(coco['images'])} images and {len(coco['annotations'])} polygons to {output_json}")