| 14 | import numpy as np |
| 15 | |
| 16 | def crop(image, target, region): |
| 17 | cropped_image = F.crop(image, *region) |
| 18 | |
| 19 | target = target.copy() |
| 20 | i, j, h, w = region |
| 21 | |
| 22 | # should we do something wrt the original size? |
| 23 | target["size"] = torch.tensor([h, w]) |
| 24 | |
| 25 | fields = ["labels", "area", "iscrowd"] |
| 26 | |
| 27 | if "boxes" in target: |
| 28 | boxes = target["boxes"] |
| 29 | max_size = torch.as_tensor([w, h], dtype=torch.float32) |
| 30 | cropped_boxes = boxes - torch.as_tensor([j, i, j, i]) |
| 31 | cropped_boxes = torch.min(cropped_boxes.reshape(-1, 2, 2), max_size) |
| 32 | cropped_boxes = cropped_boxes.clamp(min=0) |
| 33 | area = (cropped_boxes[:, 1, :] - cropped_boxes[:, 0, :]).prod(dim=1) |
| 34 | target["boxes"] = cropped_boxes.reshape(-1, 4) |
| 35 | target["area"] = area |
| 36 | fields.append("boxes") |
| 37 | |
| 38 | if "masks" in target: |
| 39 | # FIXME should we update the area here if there are no boxes? |
| 40 | target['masks'] = target['masks'][:, i:i + h, j:j + w] |
| 41 | fields.append("masks") |
| 42 | |
| 43 | # remove elements for which the boxes or masks that have zero area |
| 44 | if "boxes" in target or "masks" in target: |
| 45 | # favor boxes selection when defining which elements to keep |
| 46 | # this is compatible with previous implementation |
| 47 | if "boxes" in target: |
| 48 | cropped_boxes = target['boxes'].reshape(-1, 2, 2) |
| 49 | keep = torch.all(cropped_boxes[:, 1, :] > cropped_boxes[:, 0, :], dim=1) |
| 50 | else: |
| 51 | keep = target['masks'].flatten(1).any(1) |
| 52 | |
| 53 | for field in fields: |
| 54 | target[field] = target[field][keep] |
| 55 | |
| 56 | return cropped_image, target |
| 57 | |
| 58 | |
| 59 | def hflip(image, target): |