| 100 | |
| 101 | |
| 102 | def random_shift(image, target, region, sizes): |
| 103 | oh, ow = sizes |
| 104 | # step 1, shift crop and re-scale image firstly |
| 105 | cropped_image = F.crop(image, *region) |
| 106 | cropped_image = F.resize(cropped_image, sizes) |
| 107 | |
| 108 | target = target.copy() |
| 109 | i, j, h, w = region |
| 110 | |
| 111 | # should we do something wrt the original size? |
| 112 | target["size"] = torch.tensor([h, w]) |
| 113 | |
| 114 | fields = ["labels", "area", "iscrowd"] |
| 115 | if 'obj_ids' in target: |
| 116 | fields.append('obj_ids') |
| 117 | |
| 118 | if "boxes" in target: |
| 119 | boxes = target["boxes"] |
| 120 | max_size = torch.as_tensor([w, h], dtype=torch.float32) |
| 121 | cropped_boxes = boxes - torch.as_tensor([j, i, j, i]) |
| 122 | |
| 123 | for i, box in enumerate(cropped_boxes): |
| 124 | l, t, r, b = box |
| 125 | if l < 0: |
| 126 | l = 0 |
| 127 | if r < 0: |
| 128 | r = 0 |
| 129 | if l > w: |
| 130 | l = w |
| 131 | if r > w: |
| 132 | r = w |
| 133 | if t < 0: |
| 134 | t = 0 |
| 135 | if b < 0: |
| 136 | b = 0 |
| 137 | if t > h: |
| 138 | t = h |
| 139 | if b > h: |
| 140 | b = h |
| 141 | # step 2, re-scale coords secondly |
| 142 | ratio_h = 1.0 * oh / h |
| 143 | ratio_w = 1.0 * ow / w |
| 144 | cropped_boxes[i] = torch.tensor([ratio_w * l, ratio_h * t, ratio_w * r, ratio_h * b], dtype=box.dtype) |
| 145 | |
| 146 | cropped_boxes = cropped_boxes.reshape(-1, 2, 2) |
| 147 | area = (cropped_boxes[:, 1, :] - cropped_boxes[:, 0, :]).prod(dim=1) |
| 148 | target["boxes"] = cropped_boxes.reshape(-1, 4) |
| 149 | target["area"] = area |
| 150 | fields.append("boxes") |
| 151 | |
| 152 | if "masks" in target: |
| 153 | # FIXME should we update the area here if there are no boxes? |
| 154 | target['masks'] = target['masks'][:, i:i + h, j:j + w] |
| 155 | fields.append("masks") |
| 156 | |
| 157 | # remove elements for which the boxes or masks that have zero area |
| 158 | if "boxes" in target or "masks" in target: |
| 159 | # favor boxes selection when defining which elements to keep |