Skip to content

Huggingface Segmentation Model

sahi.models.huggingface_segmentation

HuggingFace segmentation model wrapper for SAHI.

Supports MaskFormer, Mask2Former, and OneFormer for instance, semantic, and panoptic segmentation via Hugging Face Transformers.

Classes

HuggingfaceSegmentationModel 

HuggingfaceSegmentationModel(
    *args: Any,
    overlap_mask_area_threshold: float = 0.8,
    label_ids_to_fuse: list[int] | None = None,
    min_segment_area: int = 100,
    segmentation_type: SegmentationType = INSTANCE_SEGMENTATION,
    **kwargs: Any,
)

Bases: HuggingfaceDetectionModel

HuggingFace segmentation model.

Subclasses :class:HuggingfaceDetectionModel, reusing its processor, num_categories, token handling, and dependency checks. Supports MaskFormer, Mask2Former, and OneFormer for instance, semantic, and panoptic segmentation.

Parameters:

Name Type Description Default
overlap_mask_area_threshold
float

Overlap mask area threshold to merge or discard small disconnected parts within each binary instance mask.

 0.8
label_ids_to_fuse
list[int] | None

Label ids whose instances will be fused together (panoptic only). E.g. sky can be a single segment per image.

 None
min_segment_area
int

Segments below this contour area are dropped.

 100
segmentation_type
SegmentationType

Which segmentation head to use. Params that do not apply to the chosen type are ignored.

 INSTANCE_SEGMENTATION
Source code in sahi/models/huggingface_segmentation.py 
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def __init__(
    self,
    *args: Any,
    overlap_mask_area_threshold: float = 0.8,
    label_ids_to_fuse: list[int] | None = None,
    min_segment_area: int = 100,
    segmentation_type: SegmentationType = SegmentationType.INSTANCE_SEGMENTATION,
    **kwargs: Any,
) -> None:
    self.segmentation_type = segmentation_type
    self.overlap_mask_area_threshold = overlap_mask_area_threshold
    self.label_ids_to_fuse = label_ids_to_fuse
    self.min_segment_area = min_segment_area
    # Segmentation models default to a stricter threshold than detection (0.3).
    kwargs.setdefault("confidence_threshold", 0.5)
    super().__init__(*args, **kwargs)
Methods:
load_model 
load_model() -> None

Load model and processor from HuggingFace.

Source code in sahi/models/huggingface_segmentation.py 
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
def load_model(self) -> None:
    """Load model and processor from HuggingFace."""
    from transformers import AutoModelForUniversalSegmentation, AutoProcessor

    if self.model_path is None:
        raise ValueError("model_path must be provided for HuggingFace models")

    hf_token = os.getenv("HF_TOKEN", self._token)
    model = AutoModelForUniversalSegmentation.from_pretrained(self.model_path, token=hf_token)

    processor_kwargs: dict[str, Any] = {"use_fast": False, "token": hf_token}
    if self.image_size is not None:
        processor_kwargs |= {
            "size": {"height": self.image_size, "width": self.image_size},
            "do_resize": True,
        }
    # use_fast=True raises: AttributeError: 'SizeDict' object has no attribute 'keys'
    processor = AutoProcessor.from_pretrained(self.model_path, **processor_kwargs)

    self.set_model(model, processor)

Functions: