242 lines
7.9 KiB
Python
242 lines
7.9 KiB
Python
"""
|
|
Configuration management for VR180 streaming
|
|
"""
|
|
|
|
import yaml
|
|
from pathlib import Path
|
|
from typing import Dict, Any, List, Optional
|
|
from dataclasses import dataclass, field
|
|
|
|
|
|
@dataclass
|
|
class InputConfig:
|
|
video_path: str
|
|
start_frame: int = 0
|
|
max_frames: Optional[int] = None
|
|
|
|
|
|
@dataclass
|
|
class StreamingOptions:
|
|
mode: bool = True
|
|
buffer_frames: int = 10
|
|
write_interval: int = 1 # Write every N frames
|
|
|
|
|
|
@dataclass
|
|
class ProcessingConfig:
|
|
scale_factor: float = 0.5
|
|
adaptive_scaling: bool = True
|
|
target_gpu_usage: float = 0.7
|
|
min_scale: float = 0.25
|
|
max_scale: float = 1.0
|
|
|
|
|
|
@dataclass
|
|
class DetectionConfig:
|
|
confidence_threshold: float = 0.7
|
|
model: str = "yolov8n"
|
|
device: str = "cuda"
|
|
|
|
|
|
@dataclass
|
|
class MattingConfig:
|
|
sam2_model_cfg: str = "sam2.1_hiera_l"
|
|
sam2_checkpoint: str = "segment-anything-2/checkpoints/sam2.1_hiera_large.pt"
|
|
memory_offload: bool = True
|
|
fp16: bool = True
|
|
continuous_correction: bool = True
|
|
correction_interval: int = 300
|
|
|
|
|
|
@dataclass
|
|
class StereoConfig:
|
|
mode: str = "master_slave" # "master_slave", "independent", "joint"
|
|
master_eye: str = "left"
|
|
disparity_correction: bool = True
|
|
consistency_threshold: float = 0.3
|
|
baseline: float = 65.0 # mm
|
|
focal_length: float = 1000.0 # pixels
|
|
|
|
|
|
@dataclass
|
|
class OutputConfig:
|
|
path: str
|
|
format: str = "greenscreen" # "alpha" or "greenscreen"
|
|
background_color: List[int] = field(default_factory=lambda: [0, 255, 0])
|
|
video_codec: str = "h264_nvenc"
|
|
quality_preset: str = "p4"
|
|
crf: int = 18
|
|
maintain_sbs: bool = True
|
|
|
|
|
|
@dataclass
|
|
class HardwareConfig:
|
|
device: str = "cuda"
|
|
max_vram_gb: float = 40.0
|
|
max_ram_gb: float = 48.0
|
|
|
|
|
|
@dataclass
|
|
class RecoveryConfig:
|
|
enable_checkpoints: bool = True
|
|
checkpoint_interval: int = 1000
|
|
auto_resume: bool = True
|
|
checkpoint_dir: str = "./checkpoints"
|
|
|
|
|
|
@dataclass
|
|
class PerformanceConfig:
|
|
profile_enabled: bool = True
|
|
log_interval: int = 100
|
|
memory_monitor: bool = True
|
|
|
|
|
|
class StreamingConfig:
|
|
"""Complete configuration for VR180 streaming processing"""
|
|
|
|
def __init__(self):
|
|
self.input = InputConfig("")
|
|
self.streaming = StreamingOptions()
|
|
self.processing = ProcessingConfig()
|
|
self.detection = DetectionConfig()
|
|
self.matting = MattingConfig()
|
|
self.stereo = StereoConfig()
|
|
self.output = OutputConfig("")
|
|
self.hardware = HardwareConfig()
|
|
self.recovery = RecoveryConfig()
|
|
self.performance = PerformanceConfig()
|
|
|
|
@classmethod
|
|
def from_yaml(cls, yaml_path: str) -> 'StreamingConfig':
|
|
"""Load configuration from YAML file"""
|
|
config = cls()
|
|
|
|
with open(yaml_path, 'r') as f:
|
|
data = yaml.safe_load(f)
|
|
|
|
# Update each section
|
|
if 'input' in data:
|
|
config.input = InputConfig(**data['input'])
|
|
|
|
if 'streaming' in data:
|
|
config.streaming = StreamingOptions(**data['streaming'])
|
|
|
|
if 'processing' in data:
|
|
for key, value in data['processing'].items():
|
|
setattr(config.processing, key, value)
|
|
|
|
if 'detection' in data:
|
|
config.detection = DetectionConfig(**data['detection'])
|
|
|
|
if 'matting' in data:
|
|
config.matting = MattingConfig(**data['matting'])
|
|
|
|
if 'stereo' in data:
|
|
config.stereo = StereoConfig(**data['stereo'])
|
|
|
|
if 'output' in data:
|
|
config.output = OutputConfig(**data['output'])
|
|
|
|
if 'hardware' in data:
|
|
config.hardware = HardwareConfig(**data['hardware'])
|
|
|
|
if 'recovery' in data:
|
|
config.recovery = RecoveryConfig(**data['recovery'])
|
|
|
|
if 'performance' in data:
|
|
for key, value in data['performance'].items():
|
|
setattr(config.performance, key, value)
|
|
|
|
return config
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
"""Convert configuration to dictionary"""
|
|
return {
|
|
'input': {
|
|
'video_path': self.input.video_path,
|
|
'start_frame': self.input.start_frame,
|
|
'max_frames': self.input.max_frames
|
|
},
|
|
'streaming': {
|
|
'mode': self.streaming.mode,
|
|
'buffer_frames': self.streaming.buffer_frames,
|
|
'write_interval': self.streaming.write_interval
|
|
},
|
|
'processing': {
|
|
'scale_factor': self.processing.scale_factor,
|
|
'adaptive_scaling': self.processing.adaptive_scaling,
|
|
'target_gpu_usage': self.processing.target_gpu_usage,
|
|
'min_scale': self.processing.min_scale,
|
|
'max_scale': self.processing.max_scale
|
|
},
|
|
'detection': {
|
|
'confidence_threshold': self.detection.confidence_threshold,
|
|
'model': self.detection.model,
|
|
'device': self.detection.device
|
|
},
|
|
'matting': {
|
|
'sam2_model_cfg': self.matting.sam2_model_cfg,
|
|
'sam2_checkpoint': self.matting.sam2_checkpoint,
|
|
'memory_offload': self.matting.memory_offload,
|
|
'fp16': self.matting.fp16,
|
|
'continuous_correction': self.matting.continuous_correction,
|
|
'correction_interval': self.matting.correction_interval
|
|
},
|
|
'stereo': {
|
|
'mode': self.stereo.mode,
|
|
'master_eye': self.stereo.master_eye,
|
|
'disparity_correction': self.stereo.disparity_correction,
|
|
'consistency_threshold': self.stereo.consistency_threshold,
|
|
'baseline': self.stereo.baseline,
|
|
'focal_length': self.stereo.focal_length
|
|
},
|
|
'output': {
|
|
'path': self.output.path,
|
|
'format': self.output.format,
|
|
'background_color': self.output.background_color,
|
|
'video_codec': self.output.video_codec,
|
|
'quality_preset': self.output.quality_preset,
|
|
'crf': self.output.crf,
|
|
'maintain_sbs': self.output.maintain_sbs
|
|
},
|
|
'hardware': {
|
|
'device': self.hardware.device,
|
|
'max_vram_gb': self.hardware.max_vram_gb,
|
|
'max_ram_gb': self.hardware.max_ram_gb
|
|
},
|
|
'recovery': {
|
|
'enable_checkpoints': self.recovery.enable_checkpoints,
|
|
'checkpoint_interval': self.recovery.checkpoint_interval,
|
|
'auto_resume': self.recovery.auto_resume,
|
|
'checkpoint_dir': self.recovery.checkpoint_dir
|
|
},
|
|
'performance': {
|
|
'profile_enabled': self.performance.profile_enabled,
|
|
'log_interval': self.performance.log_interval,
|
|
'memory_monitor': self.performance.memory_monitor
|
|
}
|
|
}
|
|
|
|
def validate(self) -> List[str]:
|
|
"""Validate configuration and return list of errors"""
|
|
errors = []
|
|
|
|
# Check input
|
|
if not self.input.video_path:
|
|
errors.append("Input video path is required")
|
|
elif not Path(self.input.video_path).exists():
|
|
errors.append(f"Input video not found: {self.input.video_path}")
|
|
|
|
# Check output
|
|
if not self.output.path:
|
|
errors.append("Output path is required")
|
|
|
|
# Check scale factor
|
|
if not 0.1 <= self.processing.scale_factor <= 1.0:
|
|
errors.append("Scale factor must be between 0.1 and 1.0")
|
|
|
|
# Check SAM2 checkpoint
|
|
if not Path(self.matting.sam2_checkpoint).exists():
|
|
errors.append(f"SAM2 checkpoint not found: {self.matting.sam2_checkpoint}")
|
|
|
|
return errors |