streaming part1
This commit is contained in:
298
vr180_streaming/main.py
Normal file
298
vr180_streaming/main.py
Normal file
@@ -0,0 +1,298 @@
|
||||
#!/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())
|
||||
Reference in New Issue
Block a user