Skip to content

modosaic.segmentation.generators.impl.sam_2_hiera_small

modosaic.segmentation.generators.impl.sam_2_hiera_small

SAM2HieraSmall

SAM2HieraSmall()

Bases: SegmentationGenerator

SAM 2 Hiera Small mask-generation pipeline.

Load the pinned SAM 2 Transformers pipeline.

Source code in modosaic/segmentation/generators/impl/sam_2_hiera_small.py
21
22
23
24
25
26
27
28
29
30
31
def __init__(self):
    """Load the pinned SAM 2 Transformers pipeline."""
    super().__init__()
    self.model_specs = HFModelSpec(
        model_id="facebook/sam2-hiera-small",
        revision="e080ada8afd19df5e165abe71b006edc7f4c3d4e",
        # The Transformers SAM mask postprocessor runs Torchvision NMS with
        # float32 boxes, so scores must stay float32 as well.
        dtype=torch.float32,
    )
    self.pipeline = self._load_pipeline()

generate

generate(record)

Generate segmentation masks for one image record.

Source code in modosaic/segmentation/generators/impl/sam_2_hiera_small.py
43
44
45
46
47
48
49
50
51
52
53
@override
def generate(self, record: ImageRecord) -> list[np.ndarray]:
    """Generate segmentation masks for one image record."""
    logger.debug(f"Generating segmentation masks for record sample {record.sample_id}")
    img = ImageService.bytes_to_pil(record.image_bytes)
    masks = self.pipeline(img, points_per_batch=64)["masks"]
    logger.debug(f"Generated {len(masks)} segmentation masks for record sample {record.sample_id}")
    return [
        mask.detach().cpu().float().numpy()
        for mask in masks
    ]