yslan/ObjCtrl-2.5D

1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3 4# This source code is licensed under the license found in the5# LICENSE file in the root directory of this source tree.6 7# Adapted from https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/automatic_mask_generator.py8from typing import Any, Dict, List, Optional, Tuple9 10import numpy as np11import torch12from torchvision.ops.boxes import batched_nms, box_area  # type: ignore13 14from sam2.modeling.sam2_base import SAM2Base15from sam2.sam2_image_predictor import SAM2ImagePredictor16from sam2.utils.amg import (17    area_from_rle,18    batch_iterator,19    batched_mask_to_box,20    box_xyxy_to_xywh,21    build_all_layer_point_grids,22    calculate_stability_score,23    coco_encode_rle,24    generate_crop_boxes,25    is_box_near_crop_edge,26    mask_to_rle_pytorch,27    MaskData,28    remove_small_regions,29    rle_to_mask,30    uncrop_boxes_xyxy,31    uncrop_masks,32    uncrop_points,33)34 35 36class SAM2AutomaticMaskGenerator:37    def __init__(38        self,39        model: SAM2Base,40        points_per_side: Optional[int] = 32,41        points_per_batch: int = 64,42        pred_iou_thresh: float = 0.8,43        stability_score_thresh: float = 0.95,44        stability_score_offset: float = 1.0,45        mask_threshold: float = 0.0,46        box_nms_thresh: float = 0.7,47        crop_n_layers: int = 0,48        crop_nms_thresh: float = 0.7,49        crop_overlap_ratio: float = 512 / 1500,50        crop_n_points_downscale_factor: int = 1,51        point_grids: Optional[List[np.ndarray]] = None,52        min_mask_region_area: int = 0,53        output_mode: str = "binary_mask",54        use_m2m: bool = False,55        multimask_output: bool = True,56        **kwargs,57    ) -> None:58        """59        Using a SAM 2 model, generates masks for the entire image.60        Generates a grid of point prompts over the image, then filters61        low quality and duplicate masks. The default settings are chosen62        for SAM 2 with a HieraL backbone.63 64        Arguments:65          model (Sam): The SAM 2 model to use for mask prediction.66          points_per_side (int or None): The number of points to be sampled67            along one side of the image. The total number of points is68            points_per_side**2. If None, 'point_grids' must provide explicit69            point sampling.70          points_per_batch (int): Sets the number of points run simultaneously71            by the model. Higher numbers may be faster but use more GPU memory.72          pred_iou_thresh (float): A filtering threshold in [0,1], using the73            model's predicted mask quality.74          stability_score_thresh (float): A filtering threshold in [0,1], using75            the stability of the mask under changes to the cutoff used to binarize76            the model's mask predictions.77          stability_score_offset (float): The amount to shift the cutoff when78            calculated the stability score.79          mask_threshold (float): Threshold for binarizing the mask logits80          box_nms_thresh (float): The box IoU cutoff used by non-maximal81            suppression to filter duplicate masks.82          crop_n_layers (int): If >0, mask prediction will be run again on83            crops of the image. Sets the number of layers to run, where each84            layer has 2**i_layer number of image crops.85          crop_nms_thresh (float): The box IoU cutoff used by non-maximal86            suppression to filter duplicate masks between different crops.87          crop_overlap_ratio (float): Sets the degree to which crops overlap.88            In the first crop layer, crops will overlap by this fraction of89            the image length. Later layers with more crops scale down this overlap.90          crop_n_points_downscale_factor (int): The number of points-per-side91            sampled in layer n is scaled down by crop_n_points_downscale_factor**n.92          point_grids (list(np.ndarray) or None): A list over explicit grids93            of points used for sampling, normalized to [0,1]. The nth grid in the94            list is used in the nth crop layer. Exclusive with points_per_side.95          min_mask_region_area (int): If >0, postprocessing will be applied96            to remove disconnected regions and holes in masks with area smaller97            than min_mask_region_area. Requires opencv.98          output_mode (str): The form masks are returned in. Can be 'binary_mask',99            'uncompressed_rle', or 'coco_rle'. 'coco_rle' requires pycocotools.100            For large resolutions, 'binary_mask' may consume large amounts of101            memory.102          use_m2m (bool): Whether to add a one step refinement using previous mask predictions.103          multimask_output (bool): Whether to output multimask at each point of the grid.104        """105 106        assert (points_per_side is None) != (107            point_grids is None108        ), "Exactly one of points_per_side or point_grid must be provided."109        if points_per_side is not None:110            self.point_grids = build_all_layer_point_grids(111                points_per_side,112                crop_n_layers,113                crop_n_points_downscale_factor,114            )115        elif point_grids is not None:116            self.point_grids = point_grids117        else:118            raise ValueError("Can't have both points_per_side and point_grid be None.")119 120        assert output_mode in [121            "binary_mask",122            "uncompressed_rle",123            "coco_rle",124        ], f"Unknown output_mode {output_mode}."125        if output_mode == "coco_rle":126            try:127                from pycocotools import mask as mask_utils  # type: ignore  # noqa: F401128            except ImportError as e:129                print("Please install pycocotools")130                raise e131 132        self.predictor = SAM2ImagePredictor(133            model,134            max_hole_area=min_mask_region_area,135            max_sprinkle_area=min_mask_region_area,136        )137        self.points_per_batch = points_per_batch138        self.pred_iou_thresh = pred_iou_thresh139        self.stability_score_thresh = stability_score_thresh140        self.stability_score_offset = stability_score_offset141        self.mask_threshold = mask_threshold142        self.box_nms_thresh = box_nms_thresh143        self.crop_n_layers = crop_n_layers144        self.crop_nms_thresh = crop_nms_thresh145        self.crop_overlap_ratio = crop_overlap_ratio146        self.crop_n_points_downscale_factor = crop_n_points_downscale_factor147        self.min_mask_region_area = min_mask_region_area148        self.output_mode = output_mode149        self.use_m2m = use_m2m150        self.multimask_output = multimask_output151 152    @classmethod153    def from_pretrained(cls, model_id: str, **kwargs) -> "SAM2AutomaticMaskGenerator":154        """155        Load a pretrained model from the Hugging Face hub.156 157        Arguments:158          model_id (str): The Hugging Face repository ID.159          **kwargs: Additional arguments to pass to the model constructor.160 161        Returns:162          (SAM2AutomaticMaskGenerator): The loaded model.163        """164        from sam2.build_sam import build_sam2_hf165 166        sam_model = build_sam2_hf(model_id, **kwargs)167        return cls(sam_model, **kwargs)168 169    @torch.no_grad()170    def generate(self, image: np.ndarray) -> List[Dict[str, Any]]:171        """172        Generates masks for the given image.173 174        Arguments:175          image (np.ndarray): The image to generate masks for, in HWC uint8 format.176 177        Returns:178           list(dict(str, any)): A list over records for masks. Each record is179             a dict containing the following keys:180               segmentation (dict(str, any) or np.ndarray): The mask. If181                 output_mode='binary_mask', is an array of shape HW. Otherwise,182                 is a dictionary containing the RLE.183               bbox (list(float)): The box around the mask, in XYWH format.184               area (int): The area in pixels of the mask.185               predicted_iou (float): The model's own prediction of the mask's186                 quality. This is filtered by the pred_iou_thresh parameter.187               point_coords (list(list(float))): The point coordinates input188                 to the model to generate this mask.189               stability_score (float): A measure of the mask's quality. This190                 is filtered on using the stability_score_thresh parameter.191               crop_box (list(float)): The crop of the image used to generate192                 the mask, given in XYWH format.193        """194 195        # Generate masks196        mask_data = self._generate_masks(image)197 198        # Encode masks199        if self.output_mode == "coco_rle":200            mask_data["segmentations"] = [201                coco_encode_rle(rle) for rle in mask_data["rles"]202            ]203        elif self.output_mode == "binary_mask":204            mask_data["segmentations"] = [rle_to_mask(rle) for rle in mask_data["rles"]]205        else:206            mask_data["segmentations"] = mask_data["rles"]207 208        # Write mask records209        curr_anns = []210        for idx in range(len(mask_data["segmentations"])):211            ann = {212                "segmentation": mask_data["segmentations"][idx],213                "area": area_from_rle(mask_data["rles"][idx]),214                "bbox": box_xyxy_to_xywh(mask_data["boxes"][idx]).tolist(),215                "predicted_iou": mask_data["iou_preds"][idx].item(),216                "point_coords": [mask_data["points"][idx].tolist()],217                "stability_score": mask_data["stability_score"][idx].item(),218                "crop_box": box_xyxy_to_xywh(mask_data["crop_boxes"][idx]).tolist(),219            }220            curr_anns.append(ann)221 222        return curr_anns223 224    def _generate_masks(self, image: np.ndarray) -> MaskData:225        orig_size = image.shape[:2]226        crop_boxes, layer_idxs = generate_crop_boxes(227            orig_size, self.crop_n_layers, self.crop_overlap_ratio228        )229 230        # Iterate over image crops231        data = MaskData()232        for crop_box, layer_idx in zip(crop_boxes, layer_idxs):233            crop_data = self._process_crop(image, crop_box, layer_idx, orig_size)234            data.cat(crop_data)235 236        # Remove duplicate masks between crops237        if len(crop_boxes) > 1:238            # Prefer masks from smaller crops239            scores = 1 / box_area(data["crop_boxes"])240            scores = scores.to(data["boxes"].device)241            keep_by_nms = batched_nms(242                data["boxes"].float(),243                scores,244                torch.zeros_like(data["boxes"][:, 0]),  # categories245                iou_threshold=self.crop_nms_thresh,246            )247            data.filter(keep_by_nms)248        data.to_numpy()249        return data250 251    def _process_crop(252        self,253        image: np.ndarray,254        crop_box: List[int],255        crop_layer_idx: int,256        orig_size: Tuple[int, ...],257    ) -> MaskData:258        # Crop the image and calculate embeddings259        x0, y0, x1, y1 = crop_box260        cropped_im = image[y0:y1, x0:x1, :]261        cropped_im_size = cropped_im.shape[:2]262        self.predictor.set_image(cropped_im)263 264        # Get points for this crop265        points_scale = np.array(cropped_im_size)[None, ::-1]266        points_for_image = self.point_grids[crop_layer_idx] * points_scale267 268        # Generate masks for this crop in batches269        data = MaskData()270        for (points,) in batch_iterator(self.points_per_batch, points_for_image):271            batch_data = self._process_batch(272                points, cropped_im_size, crop_box, orig_size, normalize=True273            )274            data.cat(batch_data)275            del batch_data276        self.predictor.reset_predictor()277 278        # Remove duplicates within this crop.279        keep_by_nms = batched_nms(280            data["boxes"].float(),281            data["iou_preds"],282            torch.zeros_like(data["boxes"][:, 0]),  # categories283            iou_threshold=self.box_nms_thresh,284        )285        data.filter(keep_by_nms)286 287        # Return to the original image frame288        data["boxes"] = uncrop_boxes_xyxy(data["boxes"], crop_box)289        data["points"] = uncrop_points(data["points"], crop_box)290        data["crop_boxes"] = torch.tensor([crop_box for _ in range(len(data["rles"]))])291 292        return data293 294    def _process_batch(295        self,296        points: np.ndarray,297        im_size: Tuple[int, ...],298        crop_box: List[int],299        orig_size: Tuple[int, ...],300        normalize=False,301    ) -> MaskData:302        orig_h, orig_w = orig_size303 304        # Run model on this batch305        points = torch.as_tensor(306            points, dtype=torch.float32, device=self.predictor.device307        )308        in_points = self.predictor._transforms.transform_coords(309            points, normalize=normalize, orig_hw=im_size310        )311        in_labels = torch.ones(312            in_points.shape[0], dtype=torch.int, device=in_points.device313        )314        masks, iou_preds, low_res_masks = self.predictor._predict(315            in_points[:, None, :],316            in_labels[:, None],317            multimask_output=self.multimask_output,318            return_logits=True,319        )320 321        # Serialize predictions and store in MaskData322        data = MaskData(323            masks=masks.flatten(0, 1),324            iou_preds=iou_preds.flatten(0, 1),325            points=points.repeat_interleave(masks.shape[1], dim=0),326            low_res_masks=low_res_masks.flatten(0, 1),327        )328        del masks329 330        if not self.use_m2m:331            # Filter by predicted IoU332            if self.pred_iou_thresh > 0.0:333                keep_mask = data["iou_preds"] > self.pred_iou_thresh334                data.filter(keep_mask)335 336            # Calculate and filter by stability score337            data["stability_score"] = calculate_stability_score(338                data["masks"], self.mask_threshold, self.stability_score_offset339            )340            if self.stability_score_thresh > 0.0:341                keep_mask = data["stability_score"] >= self.stability_score_thresh342                data.filter(keep_mask)343        else:344            # One step refinement using previous mask predictions345            in_points = self.predictor._transforms.transform_coords(346                data["points"], normalize=normalize, orig_hw=im_size347            )348            labels = torch.ones(349                in_points.shape[0], dtype=torch.int, device=in_points.device350            )351            masks, ious = self.refine_with_m2m(352                in_points, labels, data["low_res_masks"], self.points_per_batch353            )354            data["masks"] = masks.squeeze(1)355            data["iou_preds"] = ious.squeeze(1)356 357            if self.pred_iou_thresh > 0.0:358                keep_mask = data["iou_preds"] > self.pred_iou_thresh359                data.filter(keep_mask)360 361            data["stability_score"] = calculate_stability_score(362                data["masks"], self.mask_threshold, self.stability_score_offset363            )364            if self.stability_score_thresh > 0.0:365                keep_mask = data["stability_score"] >= self.stability_score_thresh366                data.filter(keep_mask)367 368        # Threshold masks and calculate boxes369        data["masks"] = data["masks"] > self.mask_threshold370        data["boxes"] = batched_mask_to_box(data["masks"])371 372        # Filter boxes that touch crop boundaries373        keep_mask = ~is_box_near_crop_edge(374            data["boxes"], crop_box, [0, 0, orig_w, orig_h]375        )376        if not torch.all(keep_mask):377            data.filter(keep_mask)378 379        # Compress to RLE380        data["masks"] = uncrop_masks(data["masks"], crop_box, orig_h, orig_w)381        data["rles"] = mask_to_rle_pytorch(data["masks"])382        del data["masks"]383 384        return data385 386    @staticmethod387    def postprocess_small_regions(388        mask_data: MaskData, min_area: int, nms_thresh: float389    ) -> MaskData:390        """391        Removes small disconnected regions and holes in masks, then reruns392        box NMS to remove any new duplicates.393 394        Edits mask_data in place.395 396        Requires open-cv as a dependency.397        """398        if len(mask_data["rles"]) == 0:399            return mask_data400 401        # Filter small disconnected regions and holes402        new_masks = []403        scores = []404        for rle in mask_data["rles"]:405            mask = rle_to_mask(rle)406 407            mask, changed = remove_small_regions(mask, min_area, mode="holes")408            unchanged = not changed409            mask, changed = remove_small_regions(mask, min_area, mode="islands")410            unchanged = unchanged and not changed411 412            new_masks.append(torch.as_tensor(mask).unsqueeze(0))413            # Give score=0 to changed masks and score=1 to unchanged masks414            # so NMS will prefer ones that didn't need postprocessing415            scores.append(float(unchanged))416 417        # Recalculate boxes and remove any new duplicates418        masks = torch.cat(new_masks, dim=0)419        boxes = batched_mask_to_box(masks)420        keep_by_nms = batched_nms(421            boxes.float(),422            torch.as_tensor(scores),423            torch.zeros_like(boxes[:, 0]),  # categories424            iou_threshold=nms_thresh,425        )426 427        # Only recalculate RLEs for masks that have changed428        for i_mask in keep_by_nms:429            if scores[i_mask] == 0.0:430                mask_torch = masks[i_mask].unsqueeze(0)431                mask_data["rles"][i_mask] = mask_to_rle_pytorch(mask_torch)[0]432                mask_data["boxes"][i_mask] = boxes[i_mask]  # update res directly433        mask_data.filter(keep_by_nms)434 435        return mask_data436 437    def refine_with_m2m(self, points, point_labels, low_res_masks, points_per_batch):438        new_masks = []439        new_iou_preds = []440 441        for cur_points, cur_point_labels, low_res_mask in batch_iterator(442            points_per_batch, points, point_labels, low_res_masks443        ):444            best_masks, best_iou_preds, _ = self.predictor._predict(445                cur_points[:, None, :],446                cur_point_labels[:, None],447                mask_input=low_res_mask[:, None, :],448                multimask_output=False,449                return_logits=True,450            )451            new_masks.append(best_masks)452            new_iou_preds.append(best_iou_preds)453        masks = torch.cat(new_masks, dim=0)454        return masks, torch.cat(new_iou_preds, dim=0)455