| 26 | |
| 27 | |
| 28 | def crop_mot(image, target, region): |
| 29 | cropped_image = F.crop(image, *region) |
| 30 | |
| 31 | target = target.copy() |
| 32 | i, j, h, w = region |
| 33 | |
| 34 | # should we do something wrt the original size? |
| 35 | target["size"] = torch.tensor([h, w]) |
| 36 | |
| 37 | fields = ["labels", "area", "iscrowd"] |
| 38 | if 'obj_ids' in target: |
| 39 | fields.append('obj_ids') |
| 40 | |
| 41 | if "boxes" in target: |
| 42 | boxes = target["boxes"] |
| 43 | max_size = torch.as_tensor([w, h], dtype=torch.float32) |
| 44 | cropped_boxes = boxes - torch.as_tensor([j, i, j, i]) |
| 45 | |
| 46 | for i, box in enumerate(cropped_boxes): |
| 47 | l, t, r, b = box |
| 48 | # if l < 0: |
| 49 | # l = 0 |
| 50 | # if r < 0: |
| 51 | # r = 0 |
| 52 | # if l > w: |
| 53 | # l = w |
| 54 | # if r > w: |
| 55 | # r = w |
| 56 | # if t < 0: |
| 57 | # t = 0 |
| 58 | # if b < 0: |
| 59 | # b = 0 |
| 60 | # if t > h: |
| 61 | # t = h |
| 62 | # if b > h: |
| 63 | # b = h |
| 64 | if l < 0 and r < 0: |
| 65 | l = r = 0 |
| 66 | if l > w and r > w: |
| 67 | l = r = w |
| 68 | if t < 0 and b < 0: |
| 69 | t = b = 0 |
| 70 | if t > h and b > h: |
| 71 | t = b = h |
| 72 | cropped_boxes[i] = torch.tensor([l, t, r, b], dtype=box.dtype) |
| 73 | |
| 74 | cropped_boxes = torch.min(cropped_boxes.reshape(-1, 2, 2), max_size) |
| 75 | cropped_boxes = cropped_boxes.clamp(min=0) |
| 76 | area = (cropped_boxes[:, 1, :] - cropped_boxes[:, 0, :]).prod(dim=1) |
| 77 | target["boxes"] = cropped_boxes.reshape(-1, 4) |
| 78 | target["area"] = area |
| 79 | fields.append("boxes") |
| 80 | |
| 81 | if "masks" in target: |
| 82 | # FIXME should we update the area here if there are no boxes? |
| 83 | target['masks'] = target['masks'][:, i:i + h, j:j + w] |
| 84 | fields.append("masks") |
| 85 | |