#!/usr/bin/env python3 """ YOLO-Assisted Video Annotation CLI Auto-annotate videos using pretrained YOLOv9t model. Outputs clean snapshots paired with YOLO format labels. Usage: python scripts/annotate.py --config configs/annotator.yaml python scripts/annotate.py --model models/yolov9t.pt --video input/video.mp4 --output output/ """ import argparse import sys from pathlib import Path # Add parent directory to path sys.path.insert(0, str(Path(__file__).parent.parent)) from yolo_annotator import YOLOAnnotator def main(): parser = argparse.ArgumentParser( description="YOLO-Assisted Video Annotation", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Examples: # Using config file python scripts/annotate.py --config configs/annotator.yaml # Command line options python scripts/annotate.py --model models/yolov9t.pt --video input/video.mp4 # With custom settings python scripts/annotate.py \\ --model models/yolov9t.pt \\ --video input/video.mp4 \\ --output output/annotations \\ --conf 0.3 \\ --fps 2 \\ --debug """ ) # Config file parser.add_argument( '--config', '-c', type=str, help='Path to YAML config file' ) # Model settings parser.add_argument( '--model', '-m', type=str, default='yolov9t.pt', help='Path to YOLO model file (.pt)' ) parser.add_argument( '--device', type=str, default='cuda', choices=['cuda', 'cpu'], help='Device to run inference on' ) # Video settings parser.add_argument( '--video', '-v', type=str, help='Path to input video file' ) parser.add_argument( '--fps', type=float, default=2.0, help='Frames per second to sample' ) parser.add_argument( '--max-frames', type=int, default=None, help='Maximum frames to process' ) parser.add_argument( '--start', type=float, default=0, help='Start time in seconds' ) parser.add_argument( '--end', type=float, default=None, help='End time in seconds' ) # Detection settings 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( '--classes', type=int, nargs='+', default=None, help='Class IDs to detect (default: all)' ) parser.add_argument( '--min-area', type=int, default=100, help='Minimum bbox area in pixels' ) # Output settings parser.add_argument( '--output', '-o', type=str, default='output/annotations', help='Output directory' ) parser.add_argument( '--debug', '-d', action='store_true', help='Save debug visualizations' ) parser.add_argument( '--no-manifest', action='store_true', help='Skip saving manifest file' ) args = parser.parse_args() # Build config from arguments if args.config: annotator = YOLOAnnotator(config_path=args.config) else: if not args.video: parser.error("Either --config or --video is required") config = { 'model': { 'path': args.model, 'device': args.device, 'conf_threshold': args.conf, 'iou_threshold': args.iou, }, 'video': { 'source': args.video, 'sample_fps': args.fps, 'max_frames': args.max_frames, 'start_time': args.start, 'end_time': args.end, }, 'detection': { 'classes': args.classes, 'min_confidence': args.conf, 'min_area': args.min_area, }, 'output': { 'directory': args.output, 'save_snapshots': True, 'save_labels': True, 'save_debug': args.debug, 'save_manifest': not args.no_manifest, }, } annotator = YOLOAnnotator(config=config) # Run annotation print("=" * 60) print("YOLO-Assisted Video Annotation") print("=" * 60) result = annotator.process_video() print("\n" + "=" * 60) print("COMPLETE") print("=" * 60) print(f"Processed frames: {result.processed_frames}") print(f"Total detections: {result.total_detections}") print(f"Output directory: {result.output_dir}") if __name__ == '__main__': main()