Initial commit: FishServer monorepo (FishAction, FishMeasure, fish_api)
Made-with: Cursor
This commit is contained in:
zaiun xu
222
FishMeasure/process_svo_batch.py
Executable file
222
FishMeasure/process_svo_batch.py
Executable file
@@ -0,0 +1,222 @@
|
||||
#!/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()
|
||||
|
||||
Reference in New Issue
Block a user