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