242 lines
6.3 KiB
Python
242 lines
6.3 KiB
Python
#!/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
|
|
sam2_model_cfg: "sam2.1_hiera_l"
|
|
sam2_checkpoint: "segment-anything-2/checkpoints/sam2.1_hiera_large.pt"
|
|
|
|
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()) |