| 18 | |
| 19 | |
| 20 | class 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 then |
| 30 | 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 logits |
| 35 | 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 to |
| 37 | the maximum area of fill_hole_area in low_res_masks. |
| 38 | """ |
| 39 | super().__init__() |
| 40 | self.model = sam_model |
| 41 | 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 state |
| 49 | self._is_image_set = False |
| 50 | self._features = None |
| 51 | self._orig_hw = None |
| 52 | # Whether the predictor is set for single image or a batch of images |
| 53 | self._is_batch = False |
| 54 | |
| 55 | # Predictor config |
| 56 | self.mask_threshold = mask_threshold |
| 57 | |
| 58 | # Spatial dim for backbone feature maps |
| 59 | self._bb_feat_sizes = [ |
| 60 | (256, 256), |
| 61 | (128, 128), |
| 62 | (64, 64), |
| 63 | ] |
| 64 | |
| 65 | @torch.no_grad() |
| 66 | def set_image( |
| 67 | self, |
| 68 | image: Union[np.ndarray, Image], |
| 69 | ) -> None: |
| 70 | """ |
| 71 | Calculates the image embeddings for the provided image, allowing |
| 72 | masks to be predicted with the 'predict' method. |
| 73 | |
| 74 | Arguments: |
| 75 | 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 Image |
| 76 | with pixel values in [0, 255]. |
| 77 | image_format (str): The color format of the image, in ['RGB', 'BGR']. |