File size: 1,516 Bytes
fd6c4a2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
# 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() |