298 lines
8.4 KiB
Python
298 lines
8.4 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
VR180 Streaming Human Matting - Main CLI entry point
|
|
"""
|
|
|
|
import argparse
|
|
import sys
|
|
from pathlib import Path
|
|
import traceback
|
|
|
|
from .config import StreamingConfig
|
|
from .streaming_processor import VR180StreamingProcessor
|
|
|
|
|
|
def create_parser() -> argparse.ArgumentParser:
|
|
"""Create command line argument parser"""
|
|
parser = argparse.ArgumentParser(
|
|
description="VR180 Streaming Human Matting - True streaming implementation",
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
epilog="""
|
|
Examples:
|
|
# Process video with streaming
|
|
vr180-streaming config-streaming.yaml
|
|
|
|
# Process with custom output
|
|
vr180-streaming config-streaming.yaml --output /path/to/output.mp4
|
|
|
|
# Generate example config
|
|
vr180-streaming --generate-config config-streaming-example.yaml
|
|
|
|
# Process specific frame range
|
|
vr180-streaming config-streaming.yaml --start-frame 1000 --max-frames 5000
|
|
"""
|
|
)
|
|
|
|
parser.add_argument(
|
|
"config",
|
|
nargs="?",
|
|
help="Path to YAML configuration file"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--generate-config",
|
|
metavar="PATH",
|
|
help="Generate example configuration file at specified path"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--output", "-o",
|
|
metavar="PATH",
|
|
help="Override output path from config"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--scale",
|
|
type=float,
|
|
metavar="FACTOR",
|
|
help="Override scale factor (0.25, 0.5, 1.0)"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--start-frame",
|
|
type=int,
|
|
metavar="N",
|
|
help="Start processing from frame N"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--max-frames",
|
|
type=int,
|
|
metavar="N",
|
|
help="Process at most N frames"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--device",
|
|
choices=["cuda", "cpu"],
|
|
help="Override processing device"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--format",
|
|
choices=["alpha", "greenscreen"],
|
|
help="Override output format"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--no-audio",
|
|
action="store_true",
|
|
help="Don't copy audio to output"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--verbose", "-v",
|
|
action="store_true",
|
|
help="Enable verbose output"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--dry-run",
|
|
action="store_true",
|
|
help="Validate configuration without processing"
|
|
)
|
|
|
|
return parser
|
|
|
|
|
|
def generate_example_config(output_path: str) -> None:
|
|
"""Generate example configuration file"""
|
|
config_content = '''# VR180 Streaming Configuration
|
|
# For RunPod or similar cloud GPU environments
|
|
|
|
input:
|
|
video_path: "/workspace/input_video.mp4"
|
|
start_frame: 0 # Start from beginning (or resume from checkpoint)
|
|
max_frames: null # Process entire video (or set limit for testing)
|
|
|
|
streaming:
|
|
mode: true # Enable streaming mode
|
|
buffer_frames: 10 # Small lookahead buffer
|
|
write_interval: 1 # Write every frame immediately
|
|
|
|
processing:
|
|
scale_factor: 0.5 # Process at 50% resolution for 8K input
|
|
adaptive_scaling: true # Dynamically adjust based on GPU load
|
|
target_gpu_usage: 0.7 # Target 70% GPU utilization
|
|
min_scale: 0.25
|
|
max_scale: 1.0
|
|
|
|
detection:
|
|
confidence_threshold: 0.7
|
|
model: "yolov8n" # Fast model for streaming
|
|
device: "cuda"
|
|
|
|
matting:
|
|
sam2_model_cfg: "sam2.1_hiera_l" # Large model for quality
|
|
sam2_checkpoint: "segment-anything-2/checkpoints/sam2.1_hiera_large.pt"
|
|
memory_offload: true # Essential for streaming
|
|
fp16: true # Use half precision
|
|
continuous_correction: true # Refine tracking periodically
|
|
correction_interval: 300 # Every 300 frames
|
|
|
|
stereo:
|
|
mode: "master_slave" # Left eye leads, right follows
|
|
master_eye: "left"
|
|
disparity_correction: true # Adjust for stereo depth
|
|
consistency_threshold: 0.3
|
|
baseline: 65.0 # mm - typical eye separation
|
|
focal_length: 1000.0 # pixels - adjust based on camera
|
|
|
|
output:
|
|
path: "/workspace/output_video.mp4"
|
|
format: "greenscreen" # or "alpha"
|
|
background_color: [0, 255, 0] # Pure green
|
|
video_codec: "h264_nvenc" # GPU encoding
|
|
quality_preset: "p4" # Balance quality/speed
|
|
crf: 18 # High quality
|
|
maintain_sbs: true # Keep side-by-side format
|
|
|
|
hardware:
|
|
device: "cuda"
|
|
max_vram_gb: 40.0 # RunPod A6000 has 48GB
|
|
max_ram_gb: 48.0 # Container RAM limit
|
|
|
|
recovery:
|
|
enable_checkpoints: true
|
|
checkpoint_interval: 1000 # Every 1000 frames
|
|
auto_resume: true # Resume from checkpoint if found
|
|
checkpoint_dir: "./checkpoints"
|
|
|
|
performance:
|
|
profile_enabled: true
|
|
log_interval: 100 # Log every 100 frames
|
|
memory_monitor: true # Track memory usage
|
|
'''
|
|
|
|
output_path = Path(output_path)
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
with open(output_path, 'w') as f:
|
|
f.write(config_content)
|
|
|
|
print(f"✅ Generated example configuration: {output_path}")
|
|
print("\nEdit the configuration file with your paths and run:")
|
|
print(f" python -m vr180_streaming {output_path}")
|
|
|
|
|
|
def validate_config(config: StreamingConfig, verbose: bool = False) -> bool:
|
|
"""Validate configuration and print any errors"""
|
|
errors = config.validate()
|
|
|
|
if errors:
|
|
print("❌ Configuration validation failed:")
|
|
for error in errors:
|
|
print(f" - {error}")
|
|
return False
|
|
|
|
if verbose:
|
|
print("✅ Configuration validation passed")
|
|
print(f" Input: {config.input.video_path}")
|
|
print(f" Output: {config.output.path}")
|
|
print(f" Scale: {config.processing.scale_factor}")
|
|
print(f" Device: {config.hardware.device}")
|
|
print(f" Format: {config.output.format}")
|
|
|
|
return True
|
|
|
|
|
|
def apply_cli_overrides(config: StreamingConfig, args: argparse.Namespace) -> None:
|
|
"""Apply command line overrides to configuration"""
|
|
if args.output:
|
|
config.output.path = args.output
|
|
|
|
if args.scale:
|
|
if not 0.1 <= args.scale <= 1.0:
|
|
raise ValueError("Scale factor must be between 0.1 and 1.0")
|
|
config.processing.scale_factor = args.scale
|
|
|
|
if args.start_frame is not None:
|
|
if args.start_frame < 0:
|
|
raise ValueError("Start frame must be non-negative")
|
|
config.input.start_frame = args.start_frame
|
|
|
|
if args.max_frames is not None:
|
|
if args.max_frames <= 0:
|
|
raise ValueError("Max frames must be positive")
|
|
config.input.max_frames = args.max_frames
|
|
|
|
if args.device:
|
|
config.hardware.device = args.device
|
|
config.detection.device = args.device
|
|
|
|
if args.format:
|
|
config.output.format = args.format
|
|
|
|
if args.no_audio:
|
|
config.output.maintain_sbs = False # This will skip audio copy
|
|
|
|
|
|
def main() -> int:
|
|
"""Main entry point"""
|
|
parser = create_parser()
|
|
args = parser.parse_args()
|
|
|
|
try:
|
|
# Handle config generation
|
|
if args.generate_config:
|
|
generate_example_config(args.generate_config)
|
|
return 0
|
|
|
|
# Require config file for processing
|
|
if not args.config:
|
|
parser.print_help()
|
|
print("\n❌ Error: Configuration file required")
|
|
print("\nGenerate an example config with:")
|
|
print(" vr180-streaming --generate-config config-streaming.yaml")
|
|
return 1
|
|
|
|
# Load configuration
|
|
config_path = Path(args.config)
|
|
if not config_path.exists():
|
|
print(f"❌ Error: Configuration file not found: {config_path}")
|
|
return 1
|
|
|
|
print(f"📄 Loading configuration from {config_path}")
|
|
config = StreamingConfig.from_yaml(str(config_path))
|
|
|
|
# Apply CLI overrides
|
|
apply_cli_overrides(config, args)
|
|
|
|
# Validate configuration
|
|
if not validate_config(config, verbose=args.verbose):
|
|
return 1
|
|
|
|
# Dry run mode
|
|
if args.dry_run:
|
|
print("✅ Dry run completed successfully")
|
|
return 0
|
|
|
|
# Process video
|
|
processor = VR180StreamingProcessor(config)
|
|
processor.process_video()
|
|
|
|
return 0
|
|
|
|
except KeyboardInterrupt:
|
|
print("\n⚠️ Processing interrupted by user")
|
|
return 130
|
|
|
|
except Exception as e:
|
|
print(f"\n❌ Error: {e}")
|
|
if args.verbose:
|
|
traceback.print_exc()
|
|
return 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main()) |