automatic mask generation for images using samformaskgeneration this pipeline predicts binary masks for an image given an image it is a chunkpipeline because you can seperate the points in a minibatch in order to avoid oom issues use the pointsperbatch argument to control the number of points that will be processed at the same time default is 64 the pipeline works in 3 steps 1 preprocess a grid of 1024 points evenly separated is generated along with bounding boxes and point labels for more details on how the points and bounding boxes are created check the generatecropboxes function the image is also preprocessed using the imageprocessor this function yields a minibatch of pointsperbatch 2 forward feeds the outputs of preprocess to the model the image embedding is computed only once calls both self model getimageembeddings and makes sure that the gradients are not computed and the tensors and models are on the same device 3 postprocess the most important part of the automatic mask generation happens here three steps are induced imageprocessor postprocessmasks run on each minibatch loop takes in the raw output masks resizes them according to the image size and transforms there to binary masks imageprocessor filtermasks on each minibatch loop uses both prediouthresh and stabilityscores also applies a variety of filters based on non maximum suppression to remove bad masks imageprocessor postprocessmasksforamg applies the nsm on the mask to only keep relevant ones arguments model pretrainedmodel or tfpretrainedmodel the model that will be used by the pipeline to make predictions this needs to be a model inheriting from pretrainedmodel for pytorch and tfpretrainedmodel for tensorflow tokenizer pretrainedtokenizer the tokenizer that will be used by the pipeline to encode data for the model this object inherits from pretrainedtokenizer featureextractor sequencefeatureextractor the feature extractor that will be used by the pipeline to encode the input pointsperbatch optional int default to 64 sets the number of points run simultaneously by the model higher numbers may be faster but use more gpu memory outputbboxesmask bool optional default to false whether or not to output the bounding box predictions outputrlemasks bool optional default to false whether or not to output the masks in rle format example python from transformers import pipeline generator pipelinemodelfacebooksamvitbase taskmaskgeneration outputs generator http images cocodataset orgval2017000000039769 jpg outputs generator https huggingface codatasetsnarsilimagedummyrawmainparrots png pointsperbatch128 learn more about the basics of using a pipeline in the pipeline tutorial pipelinetutorial this segmentation pipeline can currently be loaded from pipeline using the following task identifier maskgeneration see the list of available models on huggingface comodelshttps huggingface comodels filtermaskgeneration preprocess args postprocess args generates binary segmentation masks args inputs np ndarray or bytes or str or dict image or list of images maskthreshold float optional defaults to 0 0 threshold to use when turning the predicted masks into binary values prediouthresh float optional defaults to 0 88 a filtering threshold in 0 1 applied on the model s predicted mask quality stabilityscorethresh float optional defaults to 0 95 a filtering threshold in 0 1 using the stability of the mask under changes to the cutoff used to binarize the model s mask predictions stabilityscoreoffset int optional defaults to 1 the amount to shift the cutoff when calculated the stability score cropsnmsthresh float optional defaults to 0 7 the box iou cutoff used by nonmaximal suppression to filter duplicate masks cropsnlayers int optional defaults to 0 if cropsnlayers0 mask prediction will be run again on crops of the image sets the number of layers to run where each layer has 2ilayer number of image crops cropoverlapratio float optional defaults to 512 1500 sets the degree to which crops overlap in the first crop layer crops will overlap by this fraction of the image length later layers with more crops scale down this overlap cropnpointsdownscalefactor int optional defaults to 1 the number of pointsperside sampled in layer n is scaled down by cropnpointsdownscalefactorn timeout float optional defaults to none the maximum time in seconds to wait for fetching images from the web if none no timeout is set and the call may block forever return dict a dictionary with the following keys mask pil image a binary mask of the detected object as a pil image of shape width height of the original image returns a mask filled with zeros if no object is found score optional float optionally when the model is capable of estimating a confidence of the object described by the label and the mask post processing happens here in order to avoid cpu gpu copies of all the masks automatic mask generation for images using samformaskgeneration this pipeline predicts binary masks for an image given an image it is a chunkpipeline because you can seperate the points in a mini batch in order to avoid oom issues use the points_per_batch argument to control the number of points that will be processed at the same time default is 64 the pipeline works in 3 steps 1 preprocess a grid of 1024 points evenly separated is generated along with bounding boxes and point labels for more details on how the points and bounding boxes are created check the _generate_crop_boxes function the image is also preprocessed using the image_processor this function yields a minibatch of points_per_batch 2 forward feeds the outputs of preprocess to the model the image embedding is computed only once calls both self model get_image_embeddings and makes sure that the gradients are not computed and the tensors and models are on the same device 3 postprocess the most important part of the automatic mask generation happens here three steps are induced image_processor postprocess_masks run on each minibatch loop takes in the raw output masks resizes them according to the image size and transforms there to binary masks image_processor filter_masks on each minibatch loop uses both pred_iou_thresh and stability_scores also applies a variety of filters based on non maximum suppression to remove bad masks image_processor postprocess_masks_for_amg applies the nsm on the mask to only keep relevant ones arguments model pretrainedmodel or tfpretrainedmodel the model that will be used by the pipeline to make predictions this needs to be a model inheriting from pretrainedmodel for pytorch and tfpretrainedmodel for tensorflow tokenizer pretrainedtokenizer the tokenizer that will be used by the pipeline to encode data for the model this object inherits from pretrainedtokenizer feature_extractor sequencefeatureextractor the feature extractor that will be used by the pipeline to encode the input points_per_batch optional int default to 64 sets the number of points run simultaneously by the model higher numbers may be faster but use more gpu memory output_bboxes_mask bool optional default to false whether or not to output the bounding box predictions output_rle_masks bool optional default to false whether or not to output the masks in rle format example python from transformers import pipeline generator pipeline model facebook sam vit base task mask generation outputs generator http images cocodataset org val2017 000000039769 jpg outputs generator https huggingface co datasets narsil image_dummy raw main parrots png points_per_batch 128 learn more about the basics of using a pipeline in the pipeline tutorial pipeline_tutorial this segmentation pipeline can currently be loaded from pipeline using the following task identifier mask generation see the list of available models on huggingface co models https huggingface co models filter mask generation preprocess args postprocess args generates binary segmentation masks args inputs np ndarray or bytes or str or dict image or list of images mask_threshold float optional defaults to 0 0 threshold to use when turning the predicted masks into binary values pred_iou_thresh float optional defaults to 0 88 a filtering threshold in 0 1 applied on the model s predicted mask quality stability_score_thresh float optional defaults to 0 95 a filtering threshold in 0 1 using the stability of the mask under changes to the cutoff used to binarize the model s mask predictions stability_score_offset int optional defaults to 1 the amount to shift the cutoff when calculated the stability score crops_nms_thresh float optional defaults to 0 7 the box iou cutoff used by non maximal suppression to filter duplicate masks crops_n_layers int optional defaults to 0 if crops_n_layers 0 mask prediction will be run again on crops of the image sets the number of layers to run where each layer has 2 i_layer number of image crops crop_overlap_ratio float optional defaults to 512 1500 sets the degree to which crops overlap in the first crop layer crops will overlap by this fraction of the image length later layers with more crops scale down this overlap crop_n_points_downscale_factor int optional defaults to 1 the number of points per side sampled in layer n is scaled down by crop_n_points_downscale_factor n timeout float optional defaults to none the maximum time in seconds to wait for fetching images from the web if none no timeout is set and the call may block forever return dict a dictionary with the following keys mask pil image a binary mask of the detected object as a pil image of shape width height of the original image returns a mask filled with zeros if no object is found score optional float optionally when the model is capable of estimating a confidence of the object described by the label and the mask post processing happens here in order to avoid cpu gpu copies of all the masks
from collections import defaultdict
from typing import Optional

from ..image_utils import load_image
from ..utils import (
    add_end_docstrings,
    is_torch_available,
    logging,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, ChunkPipeline


if is_torch_available():
    import torch

    from ..models.auto.modeling_auto import MODEL_FOR_MASK_GENERATION_MAPPING_NAMES

logger = logging.get_logger(__name__)


@add_end_docstrings(PIPELINE_INIT_ARGS)
class MaskGenerationPipeline(ChunkPipeline):
    

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        requires_backends(self, "vision")
        requires_backends(self, "torch")

        if self.framework != "pt":
            raise ValueError(f"The {self.__class__} is only available in PyTorch.")

        self.check_model_type(MODEL_FOR_MASK_GENERATION_MAPPING_NAMES)

    def _sanitize_parameters(self, **kwargs):
        preprocess_kwargs = {}
        postprocess_kwargs = {}
        forward_params = {}
        
        if "points_per_batch" in kwargs:
            preprocess_kwargs["points_per_batch"] = kwargs["points_per_batch"]
        if "points_per_crop" in kwargs:
            preprocess_kwargs["points_per_crop"] = kwargs["points_per_crop"]
        if "crops_n_layers" in kwargs:
            preprocess_kwargs["crops_n_layers"] = kwargs["crops_n_layers"]
        if "crop_overlap_ratio" in kwargs:
            preprocess_kwargs["crop_overlap_ratio"] = kwargs["crop_overlap_ratio"]
        if "crop_n_points_downscale_factor" in kwargs:
            preprocess_kwargs["crop_n_points_downscale_factor"] = kwargs["crop_n_points_downscale_factor"]
        if "timeout" in kwargs:
            preprocess_kwargs["timeout"] = kwargs["timeout"]
        
        if "pred_iou_thresh" in kwargs:
            forward_params["pred_iou_thresh"] = kwargs["pred_iou_thresh"]
        if "stability_score_offset" in kwargs:
            forward_params["stability_score_offset"] = kwargs["stability_score_offset"]
        if "mask_threshold" in kwargs:
            forward_params["mask_threshold"] = kwargs["mask_threshold"]
        if "stability_score_thresh" in kwargs:
            forward_params["stability_score_thresh"] = kwargs["stability_score_thresh"]
        if "crops_nms_thresh" in kwargs:
            postprocess_kwargs["crops_nms_thresh"] = kwargs["crops_nms_thresh"]
        if "output_rle_mask" in kwargs:
            postprocess_kwargs["output_rle_mask"] = kwargs["output_rle_mask"]
        if "output_bboxes_mask" in kwargs:
            postprocess_kwargs["output_bboxes_mask"] = kwargs["output_bboxes_mask"]
        return preprocess_kwargs, forward_params, postprocess_kwargs

    def __call__(self, image, *args, num_workers=None, batch_size=None, **kwargs):
        
        return super().__call__(image, *args, num_workers=num_workers, batch_size=batch_size, **kwargs)

    def preprocess(
        self,
        image,
        points_per_batch=64,
        crops_n_layers: int = 0,
        crop_overlap_ratio: float = 512 / 1500,
        points_per_crop: Optional[int] = 32,
        crop_n_points_downscale_factor: Optional[int] = 1,
        timeout: Optional[float] = None,
    ):
        image = load_image(image, timeout=timeout)
        target_size = self.image_processor.size["longest_edge"]
        crop_boxes, grid_points, cropped_images, input_labels = self.image_processor.generate_crop_boxes(
            image, target_size, crops_n_layers, crop_overlap_ratio, points_per_crop, crop_n_points_downscale_factor
        )
        model_inputs = self.image_processor(images=cropped_images, return_tensors="pt")

        with self.device_placement():
            if self.framework == "pt":
                inference_context = self.get_inference_context()
                with inference_context():
                    model_inputs = self._ensure_tensor_on_device(model_inputs, device=self.device)
                    image_embeddings = self.model.get_image_embeddings(model_inputs.pop("pixel_values"))
                    model_inputs["image_embeddings"] = image_embeddings

        n_points = grid_points.shape[1]
        points_per_batch = points_per_batch if points_per_batch is not None else n_points

        if points_per_batch <= 0:
            raise ValueError(
                "Cannot have points_per_batch<=0. Must be >=1 to returned batched outputs. "
                "To return all points at once, set points_per_batch to None"
            )

        for i in range(0, n_points, points_per_batch):
            batched_points = grid_points[:, i : i + points_per_batch, :, :]
            labels = input_labels[:, i : i + points_per_batch]
            is_last = i == n_points - points_per_batch
            yield {
                "input_points": batched_points,
                "input_labels": labels,
                "input_boxes": crop_boxes,
                "is_last": is_last,
                **model_inputs,
            }

    def _forward(
        self,
        model_inputs,
        pred_iou_thresh=0.88,
        stability_score_thresh=0.95,
        mask_threshold=0,
        stability_score_offset=1,
    ):
        input_boxes = model_inputs.pop("input_boxes")
        is_last = model_inputs.pop("is_last")
        original_sizes = model_inputs.pop("original_sizes").tolist()
        reshaped_input_sizes = model_inputs.pop("reshaped_input_sizes").tolist()

        model_outputs = self.model(**model_inputs)

        
        low_resolution_masks = model_outputs["pred_masks"]
        masks = self.image_processor.post_process_masks(
            low_resolution_masks, original_sizes, reshaped_input_sizes, mask_threshold, binarize=False
        )
        iou_scores = model_outputs["iou_scores"]
        masks, iou_scores, boxes = self.image_processor.filter_masks(
            masks[0],
            iou_scores[0],
            original_sizes[0],
            input_boxes[0],
            pred_iou_thresh,
            stability_score_thresh,
            mask_threshold,
            stability_score_offset,
        )
        return {
            "masks": masks,
            "is_last": is_last,
            "boxes": boxes,
            "iou_scores": iou_scores,
        }

    def postprocess(
        self,
        model_outputs,
        output_rle_mask=False,
        output_bboxes_mask=False,
        crops_nms_thresh=0.7,
    ):
        all_scores = []
        all_masks = []
        all_boxes = []
        for model_output in model_outputs:
            all_scores.append(model_output.pop("iou_scores"))
            all_masks.extend(model_output.pop("masks"))
            all_boxes.append(model_output.pop("boxes"))

        all_scores = torch.cat(all_scores)
        all_boxes = torch.cat(all_boxes)
        output_masks, iou_scores, rle_mask, bounding_boxes = self.image_processor.post_process_for_mask_generation(
            all_masks, all_scores, all_boxes, crops_nms_thresh
        )

        extra = defaultdict(list)
        for output in model_outputs:
            for k, v in output.items():
                extra[k].append(v)

        optional = {}
        if output_rle_mask:
            optional["rle_mask"] = rle_mask

        if output_bboxes_mask:
            optional["bounding_boxes"] = bounding_boxes

        return {"masks": output_masks, "scores": iou_scores, **optional, **extra}