Transforms for SAM2.
(self, resolution, mask_threshold, max_hole_area=0.0, max_sprinkle_area=0.0)
| 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) |