| 171 | |
| 172 | |
| 173 | def crop(image, target, region): |
| 174 | cropped_image = F.crop(image, *region) |
| 175 | |
| 176 | target = target.copy() |
| 177 | i, j, h, w = region |
| 178 | |
| 179 | # should we do something wrt the original size? |
| 180 | target["size"] = torch.tensor([h, w]) |
| 181 | |
| 182 | fields = ["labels", "area", "iscrowd"] |
| 183 | if 'obj_ids' in target: |
| 184 | fields.append('obj_ids') |
| 185 | |
| 186 | if "boxes" in target: |
| 187 | boxes = target["boxes"] |
| 188 | max_size = torch.as_tensor([w, h], dtype=torch.float32) |
| 189 | cropped_boxes = boxes - torch.as_tensor([j, i, j, i]) |
| 190 | cropped_boxes = torch.min(cropped_boxes.reshape(-1, 2, 2), max_size) |
| 191 | cropped_boxes = cropped_boxes.clamp(min=0) |
| 192 | |
| 193 | area = (cropped_boxes[:, 1, :] - cropped_boxes[:, 0, :]).prod(dim=1) |
| 194 | target["boxes"] = cropped_boxes.reshape(-1, 4) |
| 195 | target["area"] = area |
| 196 | fields.append("boxes") |
| 197 | |
| 198 | if "masks" in target: |
| 199 | # FIXME should we update the area here if there are no boxes? |
| 200 | target['masks'] = target['masks'][:, i:i + h, j:j + w] |
| 201 | fields.append("masks") |
| 202 | |
| 203 | # remove elements for which the boxes or masks that have zero area |
| 204 | if "boxes" in target or "masks" in target: |
| 205 | # favor boxes selection when defining which elements to keep |
| 206 | # this is compatible with previous implementation |
| 207 | if "boxes" in target: |
| 208 | cropped_boxes = target['boxes'].reshape(-1, 2, 2) |
| 209 | keep = torch.all(cropped_boxes[:, 1, :] > cropped_boxes[:, 0, :], dim=1) |
| 210 | else: |
| 211 | keep = target['masks'].flatten(1).any(1) |
| 212 | |
| 213 | for field in fields: |
| 214 | target[field] = target[field][keep] |
| 215 | |
| 216 | return cropped_image, target |
| 217 | |
| 218 | |
| 219 | def hflip(image, target): |