Calculates the image embeddings for the provided image, allowing masks to be predicted with the 'predict' method. Arguments: 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
(
self,
image: Union[np.ndarray, Image],
)
| 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']. |
| 78 | """ |
| 79 | self.reset_predictor() |
| 80 | # Transform the image to the form expected by the model |
| 81 | if isinstance(image, np.ndarray): |
| 82 | logging.info("For numpy array image, we assume (HxWxC) format") |
| 83 | self._orig_hw = [image.shape[:2]] |
| 84 | elif isinstance(image, Image): |
| 85 | w, h = image.size |
| 86 | self._orig_hw = [(h, w)] |
| 87 | else: |
| 88 | raise NotImplementedError("Image format not supported") |
| 89 | |
| 90 | input_image = self._transforms(image) |
| 91 | input_image = input_image[None, ...].to(self.device) |
| 92 | |
| 93 | assert ( |
| 94 | len(input_image.shape) == 4 and input_image.shape[1] == 3 |
| 95 | ), f"input_image must be of size 1x3xHxW, got {input_image.shape}" |
| 96 | logging.info("Computing image embeddings for the provided image...") |
| 97 | backbone_out = self.model.forward_image(input_image) |
| 98 | _, vision_feats, _, _ = self.model._prepare_backbone_features(backbone_out) |
| 99 | # Add no_mem_embed, which is added to the lowest rest feat. map during training on videos |
| 100 | if self.model.directly_add_no_mem_embed: |
| 101 | vision_feats[-1] = vision_feats[-1] + self.model.no_mem_embed |
| 102 | |
| 103 | feats = [ |
| 104 | feat.permute(1, 2, 0).view(1, -1, *feat_size) |
| 105 | for feat, feat_size in zip(vision_feats[::-1], self._bb_feat_sizes[::-1]) |
| 106 | ][::-1] |
| 107 | self._features = {"image_embed": feats[-1], "high_res_feats": feats[:-1]} |
| 108 | self._is_image_set = True |
| 109 | logging.info("Image embeddings computed.") |
| 110 | |
| 111 | @torch.no_grad() |
| 112 | def set_image_batch( |
no test coverage detected