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 7import logging8 9from typing import List, Optional, Tuple, Union10 11import numpy as np12import torch13from PIL.Image import Image14 15from sam2.modeling.sam2_base import SAM2Base16 17from sam2.utils.transforms import SAM2Transforms18 19 20class SAM2ImagePredictor:21 def __init__(22 self,23 sam_model: SAM2Base,24 mask_threshold=0.0,25 max_hole_area=0.0,26 max_sprinkle_area=0.0,27 ) -> None:28 """29 Uses SAM-2 to calculate the image embedding for an image, and then30 allow repeated, efficient mask prediction given prompts.31 32 Arguments:33 sam_model (Sam-2): The model to use for mask prediction.34 mask_threshold (float): The threshold to use when converting mask logits35 to binary masks. Masks are thresholded at 0 by default.36 fill_hole_area (int): If fill_hole_area > 0, we fill small holes in up to37 the maximum area of fill_hole_area in low_res_masks.38 """39 super().__init__()40 self.model = sam_model41 self._transforms = SAM2Transforms(42 resolution=self.model.image_size,43 mask_threshold=mask_threshold,44 max_hole_area=max_hole_area,45 max_sprinkle_area=max_sprinkle_area,46 )47 48 # Predictor state49 self._is_image_set = False50 self._features = None51 self._orig_hw = None52 # Whether the predictor is set for single image or a batch of images53 self._is_batch = False54 55 # Predictor config56 self.mask_threshold = mask_threshold57 58 # Spatial dim for backbone feature maps59 self._bb_feat_sizes = [60 (256, 256),61 (128, 128),62 (64, 64),63 ]64 65 @classmethod66 def from_pretrained(cls, model_id: str, **kwargs) -> "SAM2ImagePredictor":67 """68 Load a pretrained model from the Hugging Face hub.69 70 Arguments:71 model_id (str): The Hugging Face repository ID.72 **kwargs: Additional arguments to pass to the model constructor.73 74 Returns:75 (SAM2ImagePredictor): The loaded model.76 """77 from sam2.build_sam import build_sam2_hf78 79 sam_model = build_sam2_hf(model_id, **kwargs)80 return cls(sam_model)81 82 @torch.no_grad()83 def set_image(84 self,85 image: Union[np.ndarray, Image],86 ) -> None:87 """88 Calculates the image embeddings for the provided image, allowing89 masks to be predicted with the 'predict' method.90 91 Arguments:92 image (np.ndarray or PIL Image): The input image to embed in RGB format. The image should be in HWC format if np.ndarray, or WHC format if PIL Image93 with pixel values in [0, 255].94 image_format (str): The color format of the image, in ['RGB', 'BGR'].95 """96 self.reset_predictor()97 # Transform the image to the form expected by the model98 if isinstance(image, np.ndarray):99 logging.info("For numpy array image, we assume (HxWxC) format")100 self._orig_hw = [image.shape[:2]]101 elif isinstance(image, Image):102 w, h = image.size103 self._orig_hw = [(h, w)]104 else:105 raise NotImplementedError("Image format not supported")106 107 input_image = self._transforms(image)108 input_image = input_image[None, ...].to(self.device)109 110 assert (111 len(input_image.shape) == 4 and input_image.shape[1] == 3112 ), f"input_image must be of size 1x3xHxW, got {input_image.shape}"113 logging.info("Computing image embeddings for the provided image...")114 backbone_out = self.model.forward_image(input_image)115 _, vision_feats, _, _ = self.model._prepare_backbone_features(backbone_out)116 # Add no_mem_embed, which is added to the lowest rest feat. map during training on videos117 if self.model.directly_add_no_mem_embed:118 vision_feats[-1] = vision_feats[-1] + self.model.no_mem_embed119 120 feats = [121 feat.permute(1, 2, 0).view(1, -1, *feat_size)122 for feat, feat_size in zip(vision_feats[::-1], self._bb_feat_sizes[::-1])123 ][::-1]124 self._features = {"image_embed": feats[-1], "high_res_feats": feats[:-1]}125 self._is_image_set = True126 logging.info("Image embeddings computed.")127 128 @torch.no_grad()129 def set_image_batch(130 self,131 image_list: List[Union[np.ndarray]],132 ) -> None:133 """134 Calculates the image embeddings for the provided image batch, allowing135 masks to be predicted with the 'predict_batch' method.136 137 Arguments:138 image_list (List[np.ndarray]): The input images to embed in RGB format. The image should be in HWC format if np.ndarray139 with pixel values in [0, 255].140 """141 self.reset_predictor()142 assert isinstance(image_list, list)143 self._orig_hw = []144 for image in image_list:145 assert isinstance(146 image, np.ndarray147 ), "Images are expected to be an np.ndarray in RGB format, and of shape HWC"148 self._orig_hw.append(image.shape[:2])149 # Transform the image to the form expected by the model150 img_batch = self._transforms.forward_batch(image_list)151 img_batch = img_batch.to(self.device)152 batch_size = img_batch.shape[0]153 assert (154 len(img_batch.shape) == 4 and img_batch.shape[1] == 3155 ), f"img_batch must be of size Bx3xHxW, got {img_batch.shape}"156 logging.info("Computing image embeddings for the provided images...")157 backbone_out = self.model.forward_image(img_batch)158 _, vision_feats, _, _ = self.model._prepare_backbone_features(backbone_out)159 # Add no_mem_embed, which is added to the lowest rest feat. map during training on videos160 if self.model.directly_add_no_mem_embed:161 vision_feats[-1] = vision_feats[-1] + self.model.no_mem_embed162 163 feats = [164 feat.permute(1, 2, 0).view(batch_size, -1, *feat_size)165 for feat, feat_size in zip(vision_feats[::-1], self._bb_feat_sizes[::-1])166 ][::-1]167 self._features = {"image_embed": feats[-1], "high_res_feats": feats[:-1]}168 self._is_image_set = True169 self._is_batch = True170 logging.info("Image embeddings computed.")171 172 def predict_batch(173 self,174 point_coords_batch: List[np.ndarray] = None,175 point_labels_batch: List[np.ndarray] = None,176 box_batch: List[np.ndarray] = None,177 mask_input_batch: List[np.ndarray] = None,178 multimask_output: bool = True,179 return_logits: bool = False,180 normalize_coords=True,181 ) -> Tuple[List[np.ndarray], List[np.ndarray], List[np.ndarray]]:182 """This function is very similar to predict(...), however it is used for batched mode, when the model is expected to generate predictions on multiple images.183 It returns a tupele of lists of masks, ious, and low_res_masks_logits.184 """185 assert self._is_batch, "This function should only be used when in batched mode"186 if not self._is_image_set:187 raise RuntimeError(188 "An image must be set with .set_image_batch(...) before mask prediction."189 )190 num_images = len(self._features["image_embed"])191 all_masks = []192 all_ious = []193 all_low_res_masks = []194 for img_idx in range(num_images):195 # Transform input prompts196 point_coords = (197 point_coords_batch[img_idx] if point_coords_batch is not None else None198 )199 point_labels = (200 point_labels_batch[img_idx] if point_labels_batch is not None else None201 )202 box = box_batch[img_idx] if box_batch is not None else None203 mask_input = (204 mask_input_batch[img_idx] if mask_input_batch is not None else None205 )206 mask_input, unnorm_coords, labels, unnorm_box = self._prep_prompts(207 point_coords,208 point_labels,209 box,210 mask_input,211 normalize_coords,212 img_idx=img_idx,213 )214 masks, iou_predictions, low_res_masks = self._predict(215 unnorm_coords,216 labels,217 unnorm_box,218 mask_input,219 multimask_output,220 return_logits=return_logits,221 img_idx=img_idx,222 )223 masks_np = masks.squeeze(0).float().detach().cpu().numpy()224 iou_predictions_np = (225 iou_predictions.squeeze(0).float().detach().cpu().numpy()226 )227 low_res_masks_np = low_res_masks.squeeze(0).float().detach().cpu().numpy()228 all_masks.append(masks_np)229 all_ious.append(iou_predictions_np)230 all_low_res_masks.append(low_res_masks_np)231 232 return all_masks, all_ious, all_low_res_masks233 234 def predict(235 self,236 point_coords: Optional[np.ndarray] = None,237 point_labels: Optional[np.ndarray] = None,238 box: Optional[np.ndarray] = None,239 mask_input: Optional[np.ndarray] = None,240 multimask_output: bool = True,241 return_logits: bool = False,242 normalize_coords=True,243 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:244 """245 Predict masks for the given input prompts, using the currently set image.246 247 Arguments:248 point_coords (np.ndarray or None): A Nx2 array of point prompts to the249 model. Each point is in (X,Y) in pixels.250 point_labels (np.ndarray or None): A length N array of labels for the251 point prompts. 1 indicates a foreground point and 0 indicates a252 background point.253 box (np.ndarray or None): A length 4 array given a box prompt to the254 model, in XYXY format.255 mask_input (np.ndarray): A low resolution mask input to the model, typically256 coming from a previous prediction iteration. Has form 1xHxW, where257 for SAM, H=W=256.258 multimask_output (bool): If true, the model will return three masks.259 For ambiguous input prompts (such as a single click), this will often260 produce better masks than a single prediction. If only a single261 mask is needed, the model's predicted quality score can be used262 to select the best mask. For non-ambiguous prompts, such as multiple263 input prompts, multimask_output=False can give better results.264 return_logits (bool): If true, returns un-thresholded masks logits265 instead of a binary mask.266 normalize_coords (bool): If true, the point coordinates will be normalized to the range [0,1] and point_coords is expected to be wrt. image dimensions.267 268 Returns:269 (np.ndarray): The output masks in CxHxW format, where C is the270 number of masks, and (H, W) is the original image size.271 (np.ndarray): An array of length C containing the model's272 predictions for the quality of each mask.273 (np.ndarray): An array of shape CxHxW, where C is the number274 of masks and H=W=256. These low resolution logits can be passed to275 a subsequent iteration as mask input.276 """277 if not self._is_image_set:278 raise RuntimeError(279 "An image must be set with .set_image(...) before mask prediction."280 )281 282 # Transform input prompts283 284 mask_input, unnorm_coords, labels, unnorm_box = self._prep_prompts(285 point_coords, point_labels, box, mask_input, normalize_coords286 )287 288 masks, iou_predictions, low_res_masks = self._predict(289 unnorm_coords,290 labels,291 unnorm_box,292 mask_input,293 multimask_output,294 return_logits=return_logits,295 )296 297 masks_np = masks.squeeze(0).float().detach().cpu().numpy()298 iou_predictions_np = iou_predictions.squeeze(0).float().detach().cpu().numpy()299 low_res_masks_np = low_res_masks.squeeze(0).float().detach().cpu().numpy()300 return masks_np, iou_predictions_np, low_res_masks_np301 302 def _prep_prompts(303 self, point_coords, point_labels, box, mask_logits, normalize_coords, img_idx=-1304 ):305 306 unnorm_coords, labels, unnorm_box, mask_input = None, None, None, None307 if point_coords is not None:308 assert (309 point_labels is not None310 ), "point_labels must be supplied if point_coords is supplied."311 point_coords = torch.as_tensor(312 point_coords, dtype=torch.float, device=self.device313 )314 unnorm_coords = self._transforms.transform_coords(315 point_coords, normalize=normalize_coords, orig_hw=self._orig_hw[img_idx]316 )317 labels = torch.as_tensor(point_labels, dtype=torch.int, device=self.device)318 if len(unnorm_coords.shape) == 2:319 unnorm_coords, labels = unnorm_coords[None, ...], labels[None, ...]320 if box is not None:321 box = torch.as_tensor(box, dtype=torch.float, device=self.device)322 unnorm_box = self._transforms.transform_boxes(323 box, normalize=normalize_coords, orig_hw=self._orig_hw[img_idx]324 ) # Bx2x2325 if mask_logits is not None:326 mask_input = torch.as_tensor(327 mask_logits, dtype=torch.float, device=self.device328 )329 if len(mask_input.shape) == 3:330 mask_input = mask_input[None, :, :, :]331 return mask_input, unnorm_coords, labels, unnorm_box332 333 @torch.no_grad()334 def _predict(335 self,336 point_coords: Optional[torch.Tensor],337 point_labels: Optional[torch.Tensor],338 boxes: Optional[torch.Tensor] = None,339 mask_input: Optional[torch.Tensor] = None,340 multimask_output: bool = True,341 return_logits: bool = False,342 img_idx: int = -1,343 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:344 """345 Predict masks for the given input prompts, using the currently set image.346 Input prompts are batched torch tensors and are expected to already be347 transformed to the input frame using SAM2Transforms.348 349 Arguments:350 point_coords (torch.Tensor or None): A BxNx2 array of point prompts to the351 model. Each point is in (X,Y) in pixels.352 point_labels (torch.Tensor or None): A BxN array of labels for the353 point prompts. 1 indicates a foreground point and 0 indicates a354 background point.355 boxes (np.ndarray or None): A Bx4 array given a box prompt to the356 model, in XYXY format.357 mask_input (np.ndarray): A low resolution mask input to the model, typically358 coming from a previous prediction iteration. Has form Bx1xHxW, where359 for SAM, H=W=256. Masks returned by a previous iteration of the360 predict method do not need further transformation.361 multimask_output (bool): If true, the model will return three masks.362 For ambiguous input prompts (such as a single click), this will often363 produce better masks than a single prediction. If only a single364 mask is needed, the model's predicted quality score can be used365 to select the best mask. For non-ambiguous prompts, such as multiple366 input prompts, multimask_output=False can give better results.367 return_logits (bool): If true, returns un-thresholded masks logits368 instead of a binary mask.369 370 Returns:371 (torch.Tensor): The output masks in BxCxHxW format, where C is the372 number of masks, and (H, W) is the original image size.373 (torch.Tensor): An array of shape BxC containing the model's374 predictions for the quality of each mask.375 (torch.Tensor): An array of shape BxCxHxW, where C is the number376 of masks and H=W=256. These low res logits can be passed to377 a subsequent iteration as mask input.378 """379 if not self._is_image_set:380 raise RuntimeError(381 "An image must be set with .set_image(...) before mask prediction."382 )383 384 if point_coords is not None:385 concat_points = (point_coords, point_labels)386 else:387 concat_points = None388 389 # Embed prompts390 if boxes is not None:391 box_coords = boxes.reshape(-1, 2, 2)392 box_labels = torch.tensor([[2, 3]], dtype=torch.int, device=boxes.device)393 box_labels = box_labels.repeat(boxes.size(0), 1)394 # we merge "boxes" and "points" into a single "concat_points" input (where395 # boxes are added at the beginning) to sam_prompt_encoder396 if concat_points is not None:397 concat_coords = torch.cat([box_coords, concat_points[0]], dim=1)398 concat_labels = torch.cat([box_labels, concat_points[1]], dim=1)399 concat_points = (concat_coords, concat_labels)400 else:401 concat_points = (box_coords, box_labels)402 403 sparse_embeddings, dense_embeddings = self.model.sam_prompt_encoder(404 points=concat_points,405 boxes=None,406 masks=mask_input,407 )408 409 # Predict masks410 batched_mode = (411 concat_points is not None and concat_points[0].shape[0] > 1412 ) # multi object prediction413 high_res_features = [414 feat_level[img_idx].unsqueeze(0)415 for feat_level in self._features["high_res_feats"]416 ]417 low_res_masks, iou_predictions, _, _ = self.model.sam_mask_decoder(418 image_embeddings=self._features["image_embed"][img_idx].unsqueeze(0),419 image_pe=self.model.sam_prompt_encoder.get_dense_pe(),420 sparse_prompt_embeddings=sparse_embeddings,421 dense_prompt_embeddings=dense_embeddings,422 multimask_output=multimask_output,423 repeat_image=batched_mode,424 high_res_features=high_res_features,425 )426 427 # Upscale the masks to the original image resolution428 masks = self._transforms.postprocess_masks(429 low_res_masks, self._orig_hw[img_idx]430 )431 low_res_masks = torch.clamp(low_res_masks, -32.0, 32.0)432 if not return_logits:433 masks = masks > self.mask_threshold434 435 return masks, iou_predictions, low_res_masks436 437 def get_image_embedding(self) -> torch.Tensor:438 """439 Returns the image embeddings for the currently set image, with440 shape 1xCxHxW, where C is the embedding dimension and (H,W) are441 the embedding spatial dimension of SAM (typically C=256, H=W=64).442 """443 if not self._is_image_set:444 raise RuntimeError(445 "An image must be set with .set_image(...) to generate an embedding."446 )447 assert (448 self._features is not None449 ), "Features must exist if an image has been set."450 return self._features["image_embed"]451 452 @property453 def device(self) -> torch.device:454 return self.model.device455 456 def reset_predictor(self) -> None:457 """458 Resets the image embeddings and other state variables.459 """460 self._is_image_set = False461 self._features = None462 self._orig_hw = None463 self._is_batch = False464