FishServer/FishMeasure/process_svo_batch.py

#!/usr/bin/env python3
"""
Batch process SVO files using ZEDReader
Processes all SVO files in an input folder and saves results to output folders
"""

import argparse
import os
import sys
from pathlib import Path
from zed_reader import ZEDReader

def process_svo_file(svo_path, output_base_dir, record=False, save_all=False, debug=False, 
                    max_depth=1000, max_y=100, min_points=500, use_yolo=False, 
                    yolo_model_path=None, show_yolo_bbox=False):
    """
    Process a single SVO file
    
    Args:
        svo_path: Path to SVO file
        output_base_dir: Base output directory
        record: Whether to record
        save_all: Whether to save all frames
        debug: Enable debug mode
        max_depth: Maximum depth value (mm)
        max_y: Maximum Y coordinate value (mm)
        min_points: Minimum points for fish detection
        use_yolo: Whether to use YOLO detector
        yolo_model_path: Path to YOLO model file
        show_yolo_bbox: Whether to show YOLO bounding boxes
    
    Returns:
        bool: True if successful, False otherwise
    """
    # Get filename without extension
    svo_filename = os.path.basename(svo_path)
    svo_name = os.path.splitext(svo_filename)[0]
    
    # Create output directory: output_base_dir/svo_name
    output_dir = os.path.join(output_base_dir, svo_name)
    
    print(f"\n{'='*60}")
    print(f"Processing: {svo_filename}")
    print(f"Output directory: {output_dir}")
    if debug:
        print(f"Debug mode: Images -> {output_dir}/images/, Point clouds -> {output_dir}/pointcloud/")
    if use_yolo:
        print(f"YOLO detector: Enabled (model: {yolo_model_path or 'default'})")
    print(f"{'='*60}")
    
    try:
        # Create ZEDReader instance with SVO file
        zed_reader = ZEDReader(svo_path=svo_path, camera_mode=False,
                              use_yolo_detector=use_yolo,
                              yolo_model_path=yolo_model_path)
        
        # Open the SVO file
        if not zed_reader.open():
            print(f"Error: Failed to open SVO file: {svo_path}")
            return False
        
        # Process frames
        # Note: show_yolo_bbox defaults to True in zed_reader.process_frames()
        # If user didn't specify --show_yolo_bbox, args.show_yolo_bbox will be False,
        # but we want to use zed_reader's default (True), so we pass True
        # If user specified --show_yolo_bbox, args.show_yolo_bbox will be True, so we pass True
        bbox_setting = True  # Always use True to match zed_reader default behavior
        
        zed_reader.process_frames(
            save_path=output_dir,
            record=record,
            save_all=save_all,
            debug=debug,
            max_depth=max_depth,
            max_y=max_y,
            min_points=min_points,
            show_yolo_bbox=bbox_setting
        )
        
        print(f"\n✓ Successfully processed: {svo_filename}")
        return True
        
    except Exception as e:
        print(f"\n✗ Error processing {svo_filename}: {str(e)}")
        import traceback
        traceback.print_exc()
        return False

def main():
    parser = argparse.ArgumentParser(
        description='Batch process SVO files using ZEDReader'
    )
    parser.add_argument(
        '--input_dir', '-i',
        type=str,
        required=True,
        help='Input directory containing SVO files'
    )
    parser.add_argument(
        '--output_dir', '-o',
        type=str,
        default='output',
        help='Base output directory (default: output)'
    )
    parser.add_argument(
        '--save_all',
        action='store_true',
        help='Save all frames (not just fish detections)'
    )
    parser.add_argument(
        '--debug',
        action='store_true',
        help='Enable debug mode: save images with bbox and point clouds when fish detected'
    )
    parser.add_argument(
        '--max_depth',
        type=float,
        default=1000,
        help='Maximum depth value (mm), default 1000'
    )
    parser.add_argument(
        '--max_y',
        type=float,
        default=150,
        help='Maximum Y coordinate value (mm), default 100'
    )
    parser.add_argument(
        '--min_points',
        type=int,
        default=500,
        help='Minimum points for fish detection, default 500'
    )
    parser.add_argument(
        '--use_yolo',
        action='store_true',
        help='Use YOLO detector for fish detection'
    )
    parser.add_argument(
        '--yolo_model',
        type=str,
        default=None,
        help='Path to YOLO model file (.pt), default uses yolov8n.pt'
    )
    parser.add_argument(
        '--show_yolo_bbox',
        action='store_true',
        help='Draw YOLO detection bounding boxes on saved images (default: True in zed_reader)'
    )
    
    args = parser.parse_args()
    
    # show_yolo_bbox defaults to True in zed_reader.process_frames()
    # If user didn't specify --show_yolo_bbox, we'll let zed_reader use its default
    # If user specified it, we pass True explicitly
    
    # Validate input directory
    if not os.path.isdir(args.input_dir):
        print(f"Error: Input directory does not exist: {args.input_dir}")
        sys.exit(1)
    
    # Create output base directory if it doesn't exist
    os.makedirs(args.output_dir, exist_ok=True)
    
    # Find all SVO files in input directory
    svo_files = []
    for file in os.listdir(args.input_dir):
        if file.lower().endswith('.svo2'):
            svo_path = os.path.join(args.input_dir, file)
            svo_files.append(svo_path)
    
    if not svo_files:
        print(f"Error: No SVO files found in: {args.input_dir}")
        sys.exit(1)
    
    print(f"Found {len(svo_files)} SVO file(s) to process")
    print(f"Input directory: {args.input_dir}")
    print(f"Output base directory: {args.output_dir}")
    print(f"Settings: max_depth={args.max_depth}mm, max_y={args.max_y}mm, min_points={args.min_points}")
    print(f"Debug mode: {args.debug}, Save all: {args.save_all}")
    print(f"YOLO detector: {args.use_yolo}, Show bbox: {args.show_yolo_bbox}")
    if args.use_yolo and args.yolo_model:
        print(f"YOLO model: {args.yolo_model}")
    if args.debug:
        print(f"Debug mode enabled: Images will be saved to output/{{svo_name}}/images/")
        print(f"                    Point clouds will be saved to output/{{svo_name}}/pointcloud/")
    
    # Process each SVO file
    successful = 0
    failed = 0
    
    for i, svo_path in enumerate(svo_files, 1):
        print(f"\n[{i}/{len(svo_files)}] Processing: {os.path.basename(svo_path)}")
        
        success = process_svo_file(
            svo_path=svo_path,
            output_base_dir=args.output_dir,
            record=False,  # Don't record when processing SVO files
            save_all=args.save_all,
            debug=args.debug,
            max_depth=args.max_depth,
            max_y=args.max_y,
            min_points=args.min_points,
            use_yolo=args.use_yolo,
            yolo_model_path=args.yolo_model,
            show_yolo_bbox=args.show_yolo_bbox
        )
        
        if success:
            successful += 1
        else:
            failed += 1
    
    # Summary
    print(f"\n{'='*60}")
    print(f"Batch processing complete!")
    print(f"  Successful: {successful}/{len(svo_files)}")
    print(f"  Failed: {failed}/{len(svo_files)}")
    print(f"{'='*60}")

if __name__ == "__main__":
    main()