diff --git a/vr180_streaming/sam2_streaming_simple.py b/vr180_streaming/sam2_streaming_simple.py index a8402e4..203832f 100644 --- a/vr180_streaming/sam2_streaming_simple.py +++ b/vr180_streaming/sam2_streaming_simple.py @@ -28,14 +28,23 @@ class SAM2StreamingProcessor: self.device = torch.device(config.get('hardware', {}).get('device', 'cuda')) # SAM2 model configuration - model_cfg = config.get('matting', {}).get('sam2_model_cfg', 'sam2.1_hiera_l') + model_cfg_name = config.get('matting', {}).get('sam2_model_cfg', 'sam2.1_hiera_l') checkpoint = config.get('matting', {}).get('sam2_checkpoint', 'segment-anything-2/checkpoints/sam2.1_hiera_large.pt') + # Map config name to Hydra path (like the examples show) + config_mapping = { + 'sam2.1_hiera_t': 'configs/sam2.1/sam2.1_hiera_t.yaml', + 'sam2.1_hiera_s': 'configs/sam2.1/sam2.1_hiera_s.yaml', + 'sam2.1_hiera_b+': 'configs/sam2.1/sam2.1_hiera_b+.yaml', + 'sam2.1_hiera_l': 'configs/sam2.1/sam2.1_hiera_l.yaml', + } + + model_cfg = config_mapping.get(model_cfg_name, model_cfg_name) + # Build predictor (simple, clean approach) - # Note: SAM2 uses Hydra to find configs automatically in sam2/configs/ self.predictor = build_sam2_video_predictor( - model_cfg, # Just the config name, Hydra finds it automatically + model_cfg, # Relative path from sam2 package checkpoint, device=self.device, vos_optimized=True # Enable VOS optimizations for speed