ZiyuG/SAM2Point
16
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 ) -> None:57 """58 Using a SAM 2 model, generates masks for the entire image.59 Generates a grid of point prompts over the image, then filters60 low quality and duplicate masks. The default settings are chosen61 for SAM 2 with a HieraL backbone.62 63 Arguments:64 model (Sam): The SAM 2 model to use for mask prediction.65 points_per_side (int or None): The number of points to be sampled66 along one side of the image. The total number of points is67 points_per_side**2. If None, 'point_grids' must provide explicit68 point sampling.69 points_per_batch (int): Sets the number of points run simultaneously70 by the model. Higher numbers may be faster but use more GPU memory.71 pred_iou_thresh (float): A filtering threshold in [0,1], using the72 model's predicted mask quality.73 stability_score_thresh (float): A filtering threshold in [0,1], using74 the stability of the mask under changes to the cutoff used to binarize75 the model's mask predictions.76 stability_score_offset (float): The amount to shift the cutoff when77 calculated the stability score.78 mask_threshold (float): Threshold for binarizing the mask logits79 box_nms_thresh (float): The box IoU cutoff used by non-maximal80 suppression to filter duplicate masks.81 crop_n_layers (int): If >0, mask prediction will be run again on82 crops of the image. Sets the number of layers to run, where each83 layer has 2**i_layer number of image crops.84 crop_nms_thresh (float): The box IoU cutoff used by non-maximal85 suppression to filter duplicate masks between different crops.86 crop_overlap_ratio (float): Sets the degree to which crops overlap.87 In the first crop layer, crops will overlap by this fraction of88 the image length. Later layers with more crops scale down this overlap.89 crop_n_points_downscale_factor (int): The number of points-per-side90 sampled in layer n is scaled down by crop_n_points_downscale_factor**n.91 point_grids (list(np.ndarray) or None): A list over explicit grids92 of points used for sampling, normalized to [0,1]. The nth grid in the93 list is used in the nth crop layer. Exclusive with points_per_side.94 min_mask_region_area (int): If >0, postprocessing will be applied95 to remove disconnected regions and holes in masks with area smaller96 than min_mask_region_area. Requires opencv.97 output_mode (str): The form masks are returned in. Can be 'binary_mask',98 'uncompressed_rle', or 'coco_rle'. 'coco_rle' requires pycocotools.99 For large resolutions, 'binary_mask' may consume large amounts of100 memory.101 use_m2m (bool): Whether to add a one step refinement using previous mask predictions.102 multimask_output (bool): Whether to output multimask at each point of the grid.103 """104 105 assert (points_per_side is None) != (106 point_grids is None107 ), "Exactly one of points_per_side or point_grid must be provided."108 if points_per_side is not None:109 self.point_grids = build_all_layer_point_grids(110 points_per_side,111 crop_n_layers,112 crop_n_points_downscale_factor,113 )114 elif point_grids is not None:115 self.point_grids = point_grids116 else:117 raise ValueError("Can't have both points_per_side and point_grid be None.")118 119 assert output_mode in [120 "binary_mask",121 "uncompressed_rle",122 "coco_rle",123 ], f"Unknown output_mode {output_mode}."124 if output_mode == "coco_rle":125 try:126 from pycocotools import mask as mask_utils # type: ignore # noqa: F401127 except ImportError as e:128 print("Please install pycocotools")129 raise e130 131 self.predictor = SAM2ImagePredictor(132 model,133 max_hole_area=min_mask_region_area,134 max_sprinkle_area=min_mask_region_area,135 )136 self.points_per_batch = points_per_batch137 self.pred_iou_thresh = pred_iou_thresh138 self.stability_score_thresh = stability_score_thresh139 self.stability_score_offset = stability_score_offset140 self.mask_threshold = mask_threshold141 self.box_nms_thresh = box_nms_thresh142 self.crop_n_layers = crop_n_layers143 self.crop_nms_thresh = crop_nms_thresh144 self.crop_overlap_ratio = crop_overlap_ratio145 self.crop_n_points_downscale_factor = crop_n_points_downscale_factor146 self.min_mask_region_area = min_mask_region_area147 self.output_mode = output_mode148 self.use_m2m = use_m2m149 self.multimask_output = multimask_output150 151 @torch.no_grad()152 def generate(self, image: np.ndarray) -> List[Dict[str, Any]]:153 """154 Generates masks for the given image.155 156 Arguments:157 image (np.ndarray): The image to generate masks for, in HWC uint8 format.158 159 Returns:160 list(dict(str, any)): A list over records for masks. Each record is161 a dict containing the following keys:162 segmentation (dict(str, any) or np.ndarray): The mask. If163 output_mode='binary_mask', is an array of shape HW. Otherwise,164 is a dictionary containing the RLE.165 bbox (list(float)): The box around the mask, in XYWH format.166 area (int): The area in pixels of the mask.167 predicted_iou (float): The model's own prediction of the mask's168 quality. This is filtered by the pred_iou_thresh parameter.169 point_coords (list(list(float))): The point coordinates input170 to the model to generate this mask.171 stability_score (float): A measure of the mask's quality. This172 is filtered on using the stability_score_thresh parameter.173 crop_box (list(float)): The crop of the image used to generate174 the mask, given in XYWH format.175 """176 177 # Generate masks178 mask_data = self._generate_masks(image)179 180 # Encode masks181 if self.output_mode == "coco_rle":182 mask_data["segmentations"] = [183 coco_encode_rle(rle) for rle in mask_data["rles"]184 ]185 elif self.output_mode == "binary_mask":186 mask_data["segmentations"] = [rle_to_mask(rle) for rle in mask_data["rles"]]187 else:188 mask_data["segmentations"] = mask_data["rles"]189 190 # Write mask records191 curr_anns = []192 for idx in range(len(mask_data["segmentations"])):193 ann = {194 "segmentation": mask_data["segmentations"][idx],195 "area": area_from_rle(mask_data["rles"][idx]),196 "bbox": box_xyxy_to_xywh(mask_data["boxes"][idx]).tolist(),197 "predicted_iou": mask_data["iou_preds"][idx].item(),198 "point_coords": [mask_data["points"][idx].tolist()],199 "stability_score": mask_data["stability_score"][idx].item(),200 "crop_box": box_xyxy_to_xywh(mask_data["crop_boxes"][idx]).tolist(),201 }202 curr_anns.append(ann)203 204 return curr_anns205 206 def _generate_masks(self, image: np.ndarray) -> MaskData:207 orig_size = image.shape[:2]208 crop_boxes, layer_idxs = generate_crop_boxes(209 orig_size, self.crop_n_layers, self.crop_overlap_ratio210 )211 212 # Iterate over image crops213 data = MaskData()214 for crop_box, layer_idx in zip(crop_boxes, layer_idxs):215 crop_data = self._process_crop(image, crop_box, layer_idx, orig_size)216 data.cat(crop_data)217 218 # Remove duplicate masks between crops219 if len(crop_boxes) > 1:220 # Prefer masks from smaller crops221 scores = 1 / box_area(data["crop_boxes"])222 scores = scores.to(data["boxes"].device)223 keep_by_nms = batched_nms(224 data["boxes"].float(),225 scores,226 torch.zeros_like(data["boxes"][:, 0]), # categories227 iou_threshold=self.crop_nms_thresh,228 )229 data.filter(keep_by_nms)230 data.to_numpy()231 return data232 233 def _process_crop(234 self,235 image: np.ndarray,236 crop_box: List[int],237 crop_layer_idx: int,238 orig_size: Tuple[int, ...],239 ) -> MaskData:240 # Crop the image and calculate embeddings241 x0, y0, x1, y1 = crop_box242 cropped_im = image[y0:y1, x0:x1, :]243 cropped_im_size = cropped_im.shape[:2]244 self.predictor.set_image(cropped_im)245 246 # Get points for this crop247 points_scale = np.array(cropped_im_size)[None, ::-1]248 points_for_image = self.point_grids[crop_layer_idx] * points_scale249 250 # Generate masks for this crop in batches251 data = MaskData()252 for (points,) in batch_iterator(self.points_per_batch, points_for_image):253 batch_data = self._process_batch(254 points, cropped_im_size, crop_box, orig_size, normalize=True255 )256 data.cat(batch_data)257 del batch_data258 self.predictor.reset_predictor()259 260 # Remove duplicates within this crop.261 keep_by_nms = batched_nms(262 data["boxes"].float(),263 data["iou_preds"],264 torch.zeros_like(data["boxes"][:, 0]), # categories265 iou_threshold=self.box_nms_thresh,266 )267 data.filter(keep_by_nms)268 269 # Return to the original image frame270 data["boxes"] = uncrop_boxes_xyxy(data["boxes"], crop_box)271 data["points"] = uncrop_points(data["points"], crop_box)272 data["crop_boxes"] = torch.tensor([crop_box for _ in range(len(data["rles"]))])273 274 return data275 276 def _process_batch(277 self,278 points: np.ndarray,279 im_size: Tuple[int, ...],280 crop_box: List[int],281 orig_size: Tuple[int, ...],282 normalize=False,283 ) -> MaskData:284 orig_h, orig_w = orig_size285 286 # Run model on this batch287 points = torch.as_tensor(points, device=self.predictor.device)288 in_points = self.predictor._transforms.transform_coords(289 points, normalize=normalize, orig_hw=im_size290 )291 in_labels = torch.ones(292 in_points.shape[0], dtype=torch.int, device=in_points.device293 )294 masks, iou_preds, low_res_masks = self.predictor._predict(295 in_points[:, None, :],296 in_labels[:, None],297 multimask_output=self.multimask_output,298 return_logits=True,299 )300 301 # Serialize predictions and store in MaskData302 data = MaskData(303 masks=masks.flatten(0, 1),304 iou_preds=iou_preds.flatten(0, 1),305 points=points.repeat_interleave(masks.shape[1], dim=0),306 low_res_masks=low_res_masks.flatten(0, 1),307 )308 del masks309 310 if not self.use_m2m:311 # Filter by predicted IoU312 if self.pred_iou_thresh > 0.0:313 keep_mask = data["iou_preds"] > self.pred_iou_thresh314 data.filter(keep_mask)315 316 # Calculate and filter by stability score317 data["stability_score"] = calculate_stability_score(318 data["masks"], self.mask_threshold, self.stability_score_offset319 )320 if self.stability_score_thresh > 0.0:321 keep_mask = data["stability_score"] >= self.stability_score_thresh322 data.filter(keep_mask)323 else:324 # One step refinement using previous mask predictions325 in_points = self.predictor._transforms.transform_coords(326 data["points"], normalize=normalize, orig_hw=im_size327 )328 labels = torch.ones(329 in_points.shape[0], dtype=torch.int, device=in_points.device330 )331 masks, ious = self.refine_with_m2m(332 in_points, labels, data["low_res_masks"], self.points_per_batch333 )334 data["masks"] = masks.squeeze(1)335 data["iou_preds"] = ious.squeeze(1)336 337 if self.pred_iou_thresh > 0.0:338 keep_mask = data["iou_preds"] > self.pred_iou_thresh339 data.filter(keep_mask)340 341 data["stability_score"] = calculate_stability_score(342 data["masks"], self.mask_threshold, self.stability_score_offset343 )344 if self.stability_score_thresh > 0.0:345 keep_mask = data["stability_score"] >= self.stability_score_thresh346 data.filter(keep_mask)347 348 # Threshold masks and calculate boxes349 data["masks"] = data["masks"] > self.mask_threshold350 data["boxes"] = batched_mask_to_box(data["masks"])351 352 # Filter boxes that touch crop boundaries353 keep_mask = ~is_box_near_crop_edge(354 data["boxes"], crop_box, [0, 0, orig_w, orig_h]355 )356 if not torch.all(keep_mask):357 data.filter(keep_mask)358 359 # Compress to RLE360 data["masks"] = uncrop_masks(data["masks"], crop_box, orig_h, orig_w)361 data["rles"] = mask_to_rle_pytorch(data["masks"])362 del data["masks"]363 364 return data365 366 @staticmethod367 def postprocess_small_regions(368 mask_data: MaskData, min_area: int, nms_thresh: float369 ) -> MaskData:370 """371 Removes small disconnected regions and holes in masks, then reruns372 box NMS to remove any new duplicates.373 374 Edits mask_data in place.375 376 Requires open-cv as a dependency.377 """378 if len(mask_data["rles"]) == 0:379 return mask_data380 381 # Filter small disconnected regions and holes382 new_masks = []383 scores = []384 for rle in mask_data["rles"]:385 mask = rle_to_mask(rle)386 387 mask, changed = remove_small_regions(mask, min_area, mode="holes")388 unchanged = not changed389 mask, changed = remove_small_regions(mask, min_area, mode="islands")390 unchanged = unchanged and not changed391 392 new_masks.append(torch.as_tensor(mask).unsqueeze(0))393 # Give score=0 to changed masks and score=1 to unchanged masks394 # so NMS will prefer ones that didn't need postprocessing395 scores.append(float(unchanged))396 397 # Recalculate boxes and remove any new duplicates398 masks = torch.cat(new_masks, dim=0)399 boxes = batched_mask_to_box(masks)400 keep_by_nms = batched_nms(401 boxes.float(),402 torch.as_tensor(scores),403 torch.zeros_like(boxes[:, 0]), # categories404 iou_threshold=nms_thresh,405 )406 407 # Only recalculate RLEs for masks that have changed408 for i_mask in keep_by_nms:409 if scores[i_mask] == 0.0:410 mask_torch = masks[i_mask].unsqueeze(0)411 mask_data["rles"][i_mask] = mask_to_rle_pytorch(mask_torch)[0]412 mask_data["boxes"][i_mask] = boxes[i_mask] # update res directly413 mask_data.filter(keep_by_nms)414 415 return mask_data416 417 def refine_with_m2m(self, points, point_labels, low_res_masks, points_per_batch):418 new_masks = []419 new_iou_preds = []420 421 for cur_points, cur_point_labels, low_res_mask in batch_iterator(422 points_per_batch, points, point_labels, low_res_masks423 ):424 best_masks, best_iou_preds, _ = self.predictor._predict(425 cur_points[:, None, :],426 cur_point_labels[:, None],427 mask_input=low_res_mask[:, None, :],428 multimask_output=False,429 return_logits=True,430 )431 new_masks.append(best_masks)432 new_iou_preds.append(best_iou_preds)433 masks = torch.cat(new_masks, dim=0)434 return masks, torch.cat(new_iou_preds, dim=0)435