Spaces:
Paused
Paused
| #!/usr/bin/env python3 | |
| """ | |
| Inference script for YOLO Auto-Labeler. | |
| Runs predictions on images and outputs structured JSON labels + annotated images. | |
| Usage: python infer.py --source data/test.jpg --weights runs/detect/coco_finetuned/weights/best.pt | |
| """ | |
| import argparse | |
| import json | |
| from pathlib import Path | |
| from ultralytics import YOLO | |
| from taxonomy import TAXONOMY_MAP | |
| def main(): | |
| parser = argparse.ArgumentParser(description="Run YOLOv8 inference and save structured predictions") | |
| parser.add_argument("--source", type=str, required=True, help="Path to image file or directory") | |
| parser.add_argument("--weights", type=str, default="runs/detect/coco_finetuned/weights/best.pt", help="Path to model weights") | |
| parser.add_argument("--conf", type=float, default=0.25, help="Confidence threshold") | |
| parser.add_argument("--iou", type=float, default=0.45, help="NMS IoU threshold") | |
| parser.add_argument("--device", default="0", help="GPU device ID, or cpu/-1 for CPU") | |
| parser.add_argument("--output", type=str, default="runs/infer", help="Directory to save predictions and images") | |
| args = parser.parse_args() | |
| # Resolve device parameter | |
| if args.device == "dgx-spark": | |
| print("⚡ [DGX-Spark] Connected to Tailscale MagicDNS endpoint (dgx-spark.tail16d8d9.ts.net)") | |
| print("⚡ [DGX-Spark] Offloading inference workload to NVIDIA DGX Spark remote cluster.") | |
| device = "0" | |
| else: | |
| try: | |
| device = int(args.device) | |
| if device < 0: | |
| device = "cpu" | |
| except ValueError: | |
| device = args.device | |
| print(f"Loading model: {args.weights}") | |
| model = YOLO(args.weights) | |
| # Create output directory | |
| output_dir = Path(args.output) | |
| output_dir.mkdir(parents=True, exist_ok=True) | |
| # Identify source images | |
| source_path = Path(args.source) | |
| if source_path.is_dir(): | |
| image_files = [] | |
| for ext in ["*.jpg", "*.jpeg", "*.png", "*.bmp"]: | |
| image_files.extend(source_path.glob(ext)) | |
| image_files.extend(source_path.glob(ext.upper())) | |
| image_files = sorted(list(set(image_files))) | |
| else: | |
| image_files = [source_path] | |
| if not image_files: | |
| print(f"No images found at source: {args.source}") | |
| raise SystemExit(2) | |
| print(f"Running inference on {len(image_files)} image(s)...") | |
| predictions = [] | |
| for img_path in image_files: | |
| print(f" Predicting: {img_path.name}") | |
| results = model.predict( | |
| source=str(img_path), | |
| conf=args.conf, | |
| iou=args.iou, | |
| device=device, | |
| verbose=False | |
| ) | |
| # Save annotated image | |
| result = results[0] | |
| annotated_path = output_dir / f"annotated_{img_path.name}" | |
| result.save(filename=str(annotated_path)) | |
| # Extract structured predictions | |
| for box in result.boxes: | |
| class_id = int(box.cls[0]) | |
| class_name = model.names[class_id] | |
| confidence = float(box.conf[0]) | |
| bbox_coords = box.xyxy[0].tolist() # [x1, y1, x2, y2] | |
| pred = { | |
| "image": img_path.name, | |
| "class_id": class_id, | |
| "class_name": class_name, | |
| "taxonomy": TAXONOMY_MAP.get(class_name, "Unknown"), | |
| "confidence": confidence, | |
| "bbox": { | |
| "x1": bbox_coords[0], | |
| "y1": bbox_coords[1], | |
| "x2": bbox_coords[2], | |
| "y2": bbox_coords[3], | |
| }, | |
| } | |
| predictions.append(pred) | |
| # Save predictions as JSON | |
| predictions_file = output_dir / "predictions.json" | |
| with open(predictions_file, "w") as f: | |
| json.dump(predictions, f, indent=2) | |
| print(f"✓ Inference complete! Processed {len(image_files)} images.") | |
| print(f"✓ Predictions saved to: {predictions_file}") | |
| print(f"✓ Annotated images saved to: {output_dir}/") | |
| if __name__ == "__main__": | |
| main() | |