105 lines
3.5 KiB
Python
105 lines
3.5 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from PIL import Image
|
|
|
|
|
|
def _detections_to_coco(
|
|
detections: Any, # supervision.Detections
|
|
image_path: Path,
|
|
image_id: int,
|
|
category_id: int = 0, # Assuming single class 'knot'
|
|
) -> dict[str, Any]:
|
|
"""Convert supervision.Detections to COCO annotation format for one image."""
|
|
annotations = []
|
|
for i in range(len(detections)):
|
|
xyxy = detections.xyxy[i].tolist()
|
|
conf = float(detections.confidence[i]) if detections.confidence is not None else 1.0
|
|
# COCO bbox: [x_min, y_min, width, height]
|
|
x_min, y_min, x_max, y_max = xyxy
|
|
width = x_max - x_min
|
|
height = y_max - y_min
|
|
bbox = [x_min, y_min, width, height]
|
|
|
|
ann = {
|
|
"id": image_id * 1000 + i, # Simple ID scheme
|
|
"image_id": image_id,
|
|
"category_id": category_id,
|
|
"bbox": bbox,
|
|
"area": width * height,
|
|
"iscrowd": 0,
|
|
"score": conf, # Optional, for model confidence
|
|
}
|
|
annotations.append(ann)
|
|
|
|
return {"annotations": annotations}
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(
|
|
description="Auto-label new images with trained RF-DETR, output COCO JSON."
|
|
)
|
|
parser.add_argument("--weights", type=Path, required=True, help="Path to trained checkpoint")
|
|
parser.add_argument("--images-dir", type=Path, required=True, help="Directory with new images")
|
|
parser.add_argument("--output-json", type=Path, required=True, help="Output COCO JSON file")
|
|
parser.add_argument("--threshold", type=float, default=0.5)
|
|
parser.add_argument("--category-name", default="knot", help="Class name for annotations")
|
|
args = parser.parse_args()
|
|
|
|
if not args.weights.exists():
|
|
raise SystemExit(f"Weights not found: {args.weights}")
|
|
if not args.images_dir.exists():
|
|
raise SystemExit(f"Images dir not found: {args.images_dir}")
|
|
|
|
from rfdetr import RFDETRBase
|
|
|
|
model = RFDETRBase(pretrain_weights=str(args.weights))
|
|
|
|
# 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 .jpg or .png images found in {args.images_dir}")
|
|
|
|
coco_data = {
|
|
"images": [],
|
|
"annotations": [],
|
|
"categories": [{"id": 0, "name": args.category_name, "supercategory": "none"}],
|
|
}
|
|
|
|
ann_id_counter = 0
|
|
for img_id, img_path in enumerate(sorted(image_paths)):
|
|
image = Image.open(img_path).convert("RGB")
|
|
detections = model.predict(image, threshold=args.threshold)
|
|
|
|
# Add image entry
|
|
width, height = image.size
|
|
coco_data["images"].append({
|
|
"id": img_id,
|
|
"file_name": img_path.name,
|
|
"width": width,
|
|
"height": height,
|
|
})
|
|
|
|
# Convert detections to annotations
|
|
ann_data = _detections_to_coco(detections, img_path, img_id, category_id=0)
|
|
for ann in ann_data["annotations"]:
|
|
ann["id"] = ann_id_counter
|
|
ann_id_counter += 1
|
|
coco_data["annotations"].append(ann)
|
|
|
|
# Write COCO JSON
|
|
with args.output_json.open("w", encoding="utf-8") as f:
|
|
json.dump(coco_data, f, indent=2)
|
|
|
|
print(f"Auto-labeled {len(image_paths)} images -> {args.output_json}")
|
|
print(f"Total annotations: {len(coco_data['annotations'])}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|