File size: 1,104 Bytes
5c36daa
b30e7a3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5c36daa
 
b30e7a3
 
 
 
 
 
 
 
 
 
 
 
 
5c36daa
 
 
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
from typing import NamedTuple, Optional, Sequence, List

import numpy as np


class SegmentationResult(NamedTuple):
    """Result from segmentation inference."""
    masks: np.ndarray  # NxHxW binary or soft masks
    scores: Optional[np.ndarray] = None  # Confidence scores
    boxes: Optional[np.ndarray] = None  # Bounding boxes (xyxy)


class Segmenter:
    """Base interface for segmentation models."""

    name: str
    supports_batch: bool = False
    max_batch_size: int = 1

    def predict(self, frame: np.ndarray, text_prompts: Optional[list] = None) -> SegmentationResult:
        """
        Run segmentation on a single frame.

        Args:
            frame: Input image as numpy array (HxWxC)
            text_prompts: Optional list of text prompts for segmentation

        Returns:
            SegmentationResult with masks and optional metadata
        """
        raise NotImplementedError

    def predict_batch(self, frames: Sequence[np.ndarray], text_prompts: Optional[list] = None) -> Sequence[SegmentationResult]:
        return [self.predict(f, text_prompts) for f in frames]