Training fails if no bounding boxes remain in patch after affine transformation

[GitMup]

Using A.Affine like so:

def get_transform(augment) -> A.Compose or ToTensorV2:
    """This is the new transform"""
    if augment:
        transform = A.Compose(
            [
                A.Affine(rotate=(-180, 180), rotate_method="ellipse", p=0.60, mode=cv2.BORDER_REFLECT_101),
                ToTensorV2(),
            ],
            bbox_params=A.BboxParams(format="pascal_voc", label_fields=["category_ids"]),
        )

    else:
        transform = ToTensorV2()

    return transform

When this transformation rotates all bounding boxes outside of the patch, leaving no boxes inside, training will result in an index error:

IndexError: tensors used as indices must be long, int, byte or bool tensors

I've fixed this one at the dataset level by having a child dataset class repeat the transform until there is an image with a bounding box:

        name, image, targets = super().__getitem__(idx)

        while targets["boxes"].size()[0] == 0:
            name, image, targets = super().__getitem__(idx)
        return name, image, targets

This could also be fixed by setting the target tensor dtype to int, but having negative samples severely degredates the perfomance on my dataset so I've done it this way.