| # 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() |