| 234 | |
| 235 | |
| 236 | def resize(image, target, size, max_size=None): |
| 237 | # size can be min_size (scalar) or (w, h) tuple |
| 238 | |
| 239 | def get_size_with_aspect_ratio(image_size, size, max_size=None): |
| 240 | w, h = image_size |
| 241 | if max_size is not None: |
| 242 | min_original_size = float(min((w, h))) |
| 243 | max_original_size = float(max((w, h))) |
| 244 | if max_original_size / min_original_size * size > max_size: |
| 245 | size = int(round(max_size * min_original_size / max_original_size)) |
| 246 | |
| 247 | if (w <= h and w == size) or (h <= w and h == size): |
| 248 | return (h, w) |
| 249 | |
| 250 | if w < h: |
| 251 | ow = size |
| 252 | oh = int(size * h / w) |
| 253 | else: |
| 254 | oh = size |
| 255 | ow = int(size * w / h) |
| 256 | |
| 257 | return (oh, ow) |
| 258 | |
| 259 | def get_size(image_size, size, max_size=None): |
| 260 | if isinstance(size, (list, tuple)): |
| 261 | return size[::-1] |
| 262 | else: |
| 263 | return get_size_with_aspect_ratio(image_size, size, max_size) |
| 264 | |
| 265 | size = get_size(image.size, size, max_size) |
| 266 | rescaled_image = F.resize(image, size) |
| 267 | |
| 268 | if target is None: |
| 269 | return rescaled_image, None |
| 270 | |
| 271 | ratios = tuple(float(s) / float(s_orig) for s, s_orig in zip(rescaled_image.size, image.size)) |
| 272 | ratio_width, ratio_height = ratios |
| 273 | |
| 274 | target = target.copy() |
| 275 | if "boxes" in target: |
| 276 | boxes = target["boxes"] |
| 277 | scaled_boxes = boxes * torch.as_tensor([ratio_width, ratio_height, ratio_width, ratio_height]) |
| 278 | target["boxes"] = scaled_boxes |
| 279 | |
| 280 | if "area" in target: |
| 281 | area = target["area"] |
| 282 | scaled_area = area * (ratio_width * ratio_height) |
| 283 | target["area"] = scaled_area |
| 284 | |
| 285 | h, w = size |
| 286 | target["size"] = torch.tensor([h, w]) |
| 287 | |
| 288 | if "masks" in target: |
| 289 | target['masks'] = interpolate( |
| 290 | target['masks'][:, None].float(), size, mode="nearest")[:, 0] > 0.5 |
| 291 | |
| 292 | return rescaled_image, target |
| 293 | |