Update prepare_annotate.py

190a9a5 verified about 1 year ago

2.26 kB

# 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}")

	# 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}")