Initial commit: Wood knot detection model and GUI

auto_label_images.py

@@ -0,0 +1,104 @@
+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())