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