# Add import at the top
from models.wrappers.matanyone_wrapper import MatAnyOneWrapper

# In your __init__ method, initialize MatAnyone
def __init__(self, ...):
    # ... existing SAM2 initialization ...
    
    # Add MatAnyone initialization
    if self.use_matanyone:
        from matanyone.inference_core import InferenceCore  # Or wherever it's located
        matanyone_core = InferenceCore(...)  # Initialize with your config
        self.matanyone = MatAnyOneWrapper(matanyone_core, device=self.device)
    
# In your process_frame or segment_frame method
def process_frame(self, frame, ...):
    # ... existing SAM2 processing ...
    
    # Add MatAnyone refinement after SAM2
    if self.use_matanyone and sam2_mask is not None:
        # Convert SAM2 output to tensor if needed
        mask_tensor = self._prepare_mask_tensor(sam2_mask)
        image_tensor = self._prepare_image_tensor(frame)
        
        # Load component masks if available
        components = None
        if self.component_paths:
            components = {
                'hair': self._load_component('hair', frame_idx),
                'edge': self._load_component('edge', frame_idx),
                # ... other components
            }
        
        # Refine with MatAnyone
        refined_mask = self.matanyone.step(
            image_tensor,
            mask_tensor,
            components=components
        )
        
        # Convert back to numpy if needed
        final_mask = refined_mask.cpu().numpy().squeeze()