| 11 | |
| 12 | |
| 13 | class SAM2Transforms(nn.Module): |
| 14 | def __init__( |
| 15 | self, resolution, mask_threshold, max_hole_area=0.0, max_sprinkle_area=0.0 |
| 16 | ): |
| 17 | """ |
| 18 | Transforms for SAM2. |
| 19 | """ |
| 20 | super().__init__() |
| 21 | self.resolution = resolution |
| 22 | self.mask_threshold = mask_threshold |
| 23 | self.max_hole_area = max_hole_area |
| 24 | self.max_sprinkle_area = max_sprinkle_area |
| 25 | self.mean = [0.485, 0.456, 0.406] |
| 26 | self.std = [0.229, 0.224, 0.225] |
| 27 | self.to_tensor = ToTensor() |
| 28 | self.transforms = torch.jit.script( |
| 29 | nn.Sequential( |
| 30 | Resize((self.resolution, self.resolution)), |
| 31 | Normalize(self.mean, self.std), |
| 32 | ) |
| 33 | ) |
| 34 | |
| 35 | def __call__(self, x): |
| 36 | x = self.to_tensor(x) |
| 37 | return self.transforms(x) |
| 38 | |
| 39 | def forward_batch(self, img_list): |
| 40 | img_batch = [self.transforms(self.to_tensor(img)) for img in img_list] |
| 41 | img_batch = torch.stack(img_batch, dim=0) |
| 42 | return img_batch |
| 43 | |
| 44 | def transform_coords( |
| 45 | self, coords: torch.Tensor, normalize=False, orig_hw=None |
| 46 | ) -> torch.Tensor: |
| 47 | """ |
| 48 | Expects a torch tensor with length 2 in the last dimension. The coordinates can be in absolute image or normalized coordinates, |
| 49 | If the coords are in absolute image coordinates, normalize should be set to True and original image size is required. |
| 50 | |
| 51 | Returns |
| 52 | Un-normalized coordinates in the range of [0, 1] which is expected by the SAM2 model. |
| 53 | """ |
| 54 | if normalize: |
| 55 | assert orig_hw is not None |
| 56 | h, w = orig_hw |
| 57 | coords = coords.clone() |
| 58 | coords[..., 0] = coords[..., 0] / w |
| 59 | coords[..., 1] = coords[..., 1] / h |
| 60 | |
| 61 | coords = coords * self.resolution # unnormalize coords |
| 62 | return coords |
| 63 | |
| 64 | def transform_boxes( |
| 65 | self, boxes: torch.Tensor, normalize=False, orig_hw=None |
| 66 | ) -> torch.Tensor: |
| 67 | """ |
| 68 | Expects a tensor of shape Bx4. The coordinates can be in absolute image or normalized coordinates, |
| 69 | if the coords are in absolute image coordinates, normalize should be set to True and original image size is required. |
| 70 | """ |