This function crops the HD images using bilinear interpolation.
| 220 | |
| 221 | |
| 222 | class CropSampler(): |
| 223 | """This function crops the HD images using bilinear interpolation.""" |
| 224 | def __init__(self, crop_size: int = 256) -> None: |
| 225 | """Uses bilinear sampling to extract square crops. |
| 226 | |
| 227 | This module expects a high resolution image as input and a bounding |
| 228 | box, described by its' center and size. It then proceeds to extract |
| 229 | a sub-image using the provided information through bilinear |
| 230 | interpolation. |
| 231 | |
| 232 | Parameters |
| 233 | ---------- |
| 234 | crop_size: int |
| 235 | The desired size for the crop. |
| 236 | """ |
| 237 | super(CropSampler, self).__init__() |
| 238 | |
| 239 | self.crop_size = crop_size |
| 240 | x = torch.arange(0, crop_size, dtype=torch.float32) / (crop_size - 1) |
| 241 | grid_y, grid_x = torch.meshgrid(x, x) |
| 242 | |
| 243 | points = torch.stack([grid_y.flatten(), grid_x.flatten()], axis=1) |
| 244 | |
| 245 | self.grid = points.unsqueeze(dim=0) |
| 246 | |
| 247 | def _sample_padded(self, full_imgs, sampling_grid): |
| 248 | """""" |
| 249 | # Get the sub-images using bilinear interpolation |
| 250 | return F.grid_sample(full_imgs, sampling_grid, align_corners=True) |
| 251 | |
| 252 | def __call__(self, full_imgs, center, bbox_size): |
| 253 | """Crops the HD images using the provided bounding boxes. |
| 254 | |
| 255 | Parameters |
| 256 | ---------- |
| 257 | full_imgs: ImageList |
| 258 | An image list structure with the full resolution images |
| 259 | center: torch.Tensor |
| 260 | A Bx2 tensor that contains the coordinates of the center of |
| 261 | the bounding box that will be cropped from the original |
| 262 | image |
| 263 | bbox_size: torch.Tensor |
| 264 | A size B tensor that contains the size of the corp |
| 265 | |
| 266 | Returns |
| 267 | ------- |
| 268 | cropped_images: torch.Tensoror |
| 269 | The images cropped from the high resolution input |
| 270 | sampling_grid: torch.Tensor |
| 271 | The grid used to sample the crops |
| 272 | """ |
| 273 | |
| 274 | batch_size, _, H, W = full_imgs.shape |
| 275 | self.grid = self.grid.to(device=full_imgs.device) |
| 276 | transforms = torch.eye(3, |
| 277 | dtype=full_imgs.dtype, |
| 278 | device=full_imgs.device).reshape( |
| 279 | 1, 3, 3).expand(batch_size, -1, |