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

#!/usr/bin/env python3
"""
Frigate-Mini-RKNN CLI

Minimal Frigate NVR fork for RKNN inference with MP4 input.
Outputs clean snapshots paired with YOLO format annotations.

Usage:
    python scripts/frigate_mini.py --config configs/frigate_mini.yaml
    python scripts/frigate_mini.py --model models/yolov9t.rknn --video input/video.mp4
"""

import argparse
import sys
from pathlib import Path

# Add parent directory to path
sys.path.insert(0, str(Path(__file__).parent.parent))

from frigate_mini import FrigateMini


def main():
    parser = argparse.ArgumentParser(
        description="Frigate-Mini RKNN - Object Detection with Annotation Export",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Using config file
  python scripts/frigate_mini.py --config configs/frigate_mini.yaml
  
  # Quick start with model and video
  python scripts/frigate_mini.py --model models/yolov9t.pt --video input/video.mp4
  
  # With RKNN model on Rockchip device
  python scripts/frigate_mini.py \\
      --model models/yolov9t.rknn \\
      --video input/video.mp4 \\
      --output output/ \\
      --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='current.pt',
        help='Path to model file (.pt, .onnx, or .rknn) (default: current.pt)'
    )
    parser.add_argument(
        '--detector-type',
        type=str,
        choices=['rknn', 'onnx', 'yolo'],
        default=None,
        help='Detector type (auto-detected from extension if not specified)'
    )
    parser.add_argument(
        '--platform',
        type=str,
        default='rk3588',
        help='RKNN target platform (rk3588, rk3568, etc.)'
    )
    
    # Video settings
    parser.add_argument(
        '--video', '-v',
        type=str,
        default='dianote.mp4',
        help='Path to input video file (default: dianote.mp4)'
    )
    parser.add_argument(
        '--fps',
        type=int,
        default=5,
        help='Processing FPS limit'
    )
    parser.add_argument(
        '--loop',
        action='store_true',
        default=True,
        help='Loop video playback'
    )
    parser.add_argument(
        '--no-loop',
        action='store_true',
        help='Disable video loop'
    )
    
    # Detection settings
    parser.add_argument(
        '--conf',
        type=float,
        default=0.25,
        help='Confidence threshold'
    )
    parser.add_argument(
        '--objects',
        type=str,
        nargs='+',
        default=['person', 'car'],
        help='Objects to track (class names)'
    )
    
    # Output settings
    parser.add_argument(
        '--output', '-o',
        type=str,
        default='output',
        help='Output directory'
    )
    parser.add_argument(
        '--debug', '-d',
        action='store_true',
        help='Enable debug mode'
    )
    parser.add_argument(
        '--log-level',
        type=str,
        choices=['debug', 'info', 'warning', 'error'],
        default='info',
        help='Logging level'
    )
    
    args = parser.parse_args()
    
    # Build config from arguments
    if args.config:
        app = FrigateMini(config_path=args.config)
    else:
        # Use defaults if not provided
        model_path = args.model if args.model else 'current.pt'
        video_path = args.video if args.video else 'dianote.mp4'
        
        # Auto-detect detector type from extension
        detector_type = args.detector_type
        if detector_type is None:
            if model_path.endswith('.rknn'):
                detector_type = 'rknn'
            elif model_path.endswith('.onnx'):
                detector_type = 'onnx'
            else:
                detector_type = 'yolo'
        
        config = {
            'debug': args.debug,
            'log_level': args.log_level,
            
            'detector': {
                'type': detector_type,
                'model_path': model_path,
                'conf_threshold': args.conf,
                'rknn': {
                    'target_platform': args.platform,
                },
            },
            
            'cameras': {
                'default': {
                    'enabled': True,
                    'source': video_path,
                    'fps': args.fps,
                    'loop': not args.no_loop,
                    'objects': {
                        'track': args.objects,
                    },
                },
            },
            
            'snapshots': {
                'enabled': True,
                'output_dir': args.output,
                'trigger': {
                    'objects': args.objects,
                    'min_score': 0.5,
                    'cooldown': 2.0,
                },
            },
            
            'annotations': {
                'enabled': True,
                'output_dir': args.output,
                'pair_with_snapshots': True,
            },
            
            'debug_output': {
                'enabled': args.debug,
                'output_dir': f"{args.output}/debug",
            },
        }
        
        app = FrigateMini(config=config)
    
    # Run
    print("=" * 60)
    print("Frigate-Mini RKNN")
    print("=" * 60)
    
    app.start()


if __name__ == '__main__':
    main()