dataset-yolo-script/sam2-cpu/scripts/annotate.py

#!/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()