129 lines
4.4 KiB
Python
129 lines
4.4 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import json
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
from PIL import Image
|
||
|
|
|
||
|
|
|
||
|
|
def coco_to_label_studio(
|
||
|
|
image_paths: list[Path],
|
||
|
|
coco_predictions_path: Path,
|
||
|
|
category_name: str = "knot",
|
||
|
|
) -> list[dict]:
|
||
|
|
"""Convert COCO predictions to Label Studio import format with pre-annotations."""
|
||
|
|
|
||
|
|
# Load COCO predictions
|
||
|
|
with coco_predictions_path.open("r", encoding="utf-8") as f:
|
||
|
|
coco_data = json.load(f)
|
||
|
|
|
||
|
|
# Build image_id -> annotations mapping
|
||
|
|
annotations_by_image = {}
|
||
|
|
for ann in coco_data.get("annotations", []):
|
||
|
|
img_id = ann["image_id"]
|
||
|
|
if img_id not in annotations_by_image:
|
||
|
|
annotations_by_image[img_id] = []
|
||
|
|
annotations_by_image[img_id].append(ann)
|
||
|
|
|
||
|
|
# Build image_id -> image info mapping
|
||
|
|
image_info = {img["id"]: img for img in coco_data.get("images", [])}
|
||
|
|
|
||
|
|
# Convert to Label Studio format
|
||
|
|
ls_tasks = []
|
||
|
|
for img_id, img_data in image_info.items():
|
||
|
|
file_name = img_data["file_name"]
|
||
|
|
width = img_data["width"]
|
||
|
|
height = img_data["height"]
|
||
|
|
|
||
|
|
# Find the actual file path
|
||
|
|
matching_paths = [p for p in image_paths if p.name == file_name]
|
||
|
|
if not matching_paths:
|
||
|
|
continue
|
||
|
|
|
||
|
|
img_path = matching_paths[0]
|
||
|
|
|
||
|
|
# Build predictions (pre-annotations)
|
||
|
|
predictions = []
|
||
|
|
for ann in annotations_by_image.get(img_id, []):
|
||
|
|
# COCO bbox: [x_min, y_min, width, height]
|
||
|
|
x, y, w, h = ann["bbox"]
|
||
|
|
|
||
|
|
# Label Studio uses percentages
|
||
|
|
x_percent = (x / width) * 100
|
||
|
|
y_percent = (y / height) * 100
|
||
|
|
w_percent = (w / width) * 100
|
||
|
|
h_percent = (h / height) * 100
|
||
|
|
|
||
|
|
predictions.append({
|
||
|
|
"value": {
|
||
|
|
"x": x_percent,
|
||
|
|
"y": y_percent,
|
||
|
|
"width": w_percent,
|
||
|
|
"height": h_percent,
|
||
|
|
"rectanglelabels": [category_name]
|
||
|
|
},
|
||
|
|
"from_name": "label",
|
||
|
|
"to_name": "image",
|
||
|
|
"type": "rectanglelabels",
|
||
|
|
"score": ann.get("score", 1.0)
|
||
|
|
})
|
||
|
|
|
||
|
|
# Create Label Studio task with predictions
|
||
|
|
task = {
|
||
|
|
"data": {
|
||
|
|
"image": f"/data/local-files/?d={img_path.absolute()}"
|
||
|
|
},
|
||
|
|
"predictions": [{
|
||
|
|
"result": predictions,
|
||
|
|
"model_version": "rfdetr-auto-label"
|
||
|
|
}]
|
||
|
|
}
|
||
|
|
|
||
|
|
ls_tasks.append(task)
|
||
|
|
|
||
|
|
return ls_tasks
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> int:
|
||
|
|
parser = argparse.ArgumentParser(
|
||
|
|
description="Convert COCO predictions to Label Studio format for review."
|
||
|
|
)
|
||
|
|
parser.add_argument("--coco-json", type=Path, required=True, help="COCO predictions JSON from auto_label_images.py")
|
||
|
|
parser.add_argument("--images-dir", type=Path, required=True, help="Directory with the images")
|
||
|
|
parser.add_argument("--output-json", type=Path, required=True, help="Output Label Studio tasks JSON")
|
||
|
|
parser.add_argument("--category-name", default="knot", help="Label name in Label Studio")
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
if not args.coco_json.exists():
|
||
|
|
raise SystemExit(f"COCO JSON not found: {args.coco_json}")
|
||
|
|
if not args.images_dir.exists():
|
||
|
|
raise SystemExit(f"Images dir not found: {args.images_dir}")
|
||
|
|
|
||
|
|
# Collect all images
|
||
|
|
image_paths = list(args.images_dir.glob("*.jpg")) + list(args.images_dir.glob("*.png"))
|
||
|
|
if not image_paths:
|
||
|
|
raise SystemExit(f"No images found in {args.images_dir}")
|
||
|
|
|
||
|
|
# Convert
|
||
|
|
ls_tasks = coco_to_label_studio(image_paths, args.coco_json, args.category_name)
|
||
|
|
|
||
|
|
# Write Label Studio import format
|
||
|
|
with args.output_json.open("w", encoding="utf-8") as f:
|
||
|
|
json.dump(ls_tasks, f, indent=2)
|
||
|
|
|
||
|
|
print(f"✓ Converted {len(ls_tasks)} tasks with predictions")
|
||
|
|
print(f"✓ Output: {args.output_json}")
|
||
|
|
print(f"\nImport into Label Studio:")
|
||
|
|
print(f" 1. Open your project")
|
||
|
|
print(f" 2. Settings → Storage → Add Local Storage")
|
||
|
|
print(f" Absolute path: {args.images_dir.absolute()}")
|
||
|
|
print(f" 3. Import → Upload Files → select {args.output_json.name}")
|
||
|
|
print(f" 4. Start reviewing predictions!")
|
||
|
|
|
||
|
|
return 0
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
raise SystemExit(main())
|