first commit
This commit is contained in:
240
vr180_matting/main.py
Normal file
240
vr180_matting/main.py
Normal file
@@ -0,0 +1,240 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
VR180 Human Matting with Det-SAM2
|
||||
Main CLI entry point
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
from pathlib import Path
|
||||
import traceback
|
||||
|
||||
from .config import VR180Config
|
||||
from .vr180_processor import VR180Processor
|
||||
|
||||
|
||||
def create_parser() -> argparse.ArgumentParser:
|
||||
"""Create command line argument parser"""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="VR180 Human Matting with Det-SAM2",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
Examples:
|
||||
# Process video with default config
|
||||
vr180-matting config.yaml
|
||||
|
||||
# Process with custom output path
|
||||
vr180-matting config.yaml --output /path/to/output.mp4
|
||||
|
||||
# Generate example config
|
||||
vr180-matting --generate-config config_example.yaml
|
||||
|
||||
# Process with different scale factor
|
||||
vr180-matting config.yaml --scale 0.25
|
||||
"""
|
||||
)
|
||||
|
||||
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(
|
||||
"--chunk-size",
|
||||
type=int,
|
||||
metavar="FRAMES",
|
||||
help="Override chunk size in frames (0 for auto)"
|
||||
)
|
||||
|
||||
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(
|
||||
"--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 = '''input:
|
||||
video_path: "path/to/input.mp4"
|
||||
|
||||
processing:
|
||||
scale_factor: 0.5 # 0.25, 0.5, 1.0
|
||||
chunk_size: 900 # frames, 0 for full video
|
||||
overlap_frames: 60 # for chunked processing
|
||||
|
||||
detection:
|
||||
confidence_threshold: 0.7
|
||||
model: "yolov8n" # yolov8n, yolov8s, yolov8m
|
||||
|
||||
matting:
|
||||
use_disparity_mapping: true
|
||||
memory_offload: true
|
||||
fp16: true
|
||||
|
||||
output:
|
||||
path: "path/to/output/"
|
||||
format: "alpha" # "alpha" or "greenscreen"
|
||||
background_color: [0, 255, 0] # for greenscreen
|
||||
maintain_sbs: true # keep side-by-side format
|
||||
|
||||
hardware:
|
||||
device: "cuda"
|
||||
max_vram_gb: 10 # RTX 3080 limit
|
||||
'''
|
||||
|
||||
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("Edit the configuration file and run:")
|
||||
print(f" vr180-matting {output_path}")
|
||||
|
||||
|
||||
def validate_config(config: VR180Config, 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 video: {config.input.video_path}")
|
||||
print(f"Output path: {config.output.path}")
|
||||
print(f"Scale factor: {config.processing.scale_factor}")
|
||||
print(f"Device: {config.hardware.device}")
|
||||
print(f"Output format: {config.output.format}")
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def apply_cli_overrides(config: VR180Config, 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.chunk_size is not None:
|
||||
if args.chunk_size < 0:
|
||||
raise ValueError("Chunk size must be non-negative")
|
||||
config.processing.chunk_size = args.chunk_size
|
||||
|
||||
if args.device:
|
||||
config.hardware.device = args.device
|
||||
|
||||
if args.format:
|
||||
config.output.format = args.format
|
||||
|
||||
|
||||
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("\nError: Configuration file required")
|
||||
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 = VR180Config.from_yaml(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
|
||||
|
||||
# Initialize processor
|
||||
print("Initializing VR180 processor...")
|
||||
processor = VR180Processor(config)
|
||||
|
||||
# Process video
|
||||
processor.process_video()
|
||||
|
||||
print("✅ Processing completed successfully!")
|
||||
return 0
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("\n⚠️ Processing interrupted by user")
|
||||
return 130
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Error: {e}")
|
||||
if args.verbose:
|
||||
traceback.print_exc()
|
||||
return 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Reference in New Issue
Block a user