I am trying to add augmentation to my data loader in torch. All of my example images have masks and none are blank. But when I add an augmentation like RandomRotation I find that some examples in my dataset then yield an image that has not labels anymore and things fall apart.
Is there a way to correctly handle this in PyTorch. Here as example of the problem?
def __getitem__(self, idx):
# load image and mask
img_path = os.path.join(self.root, self.imagefolder, self.imgs[idx])
mask_path = os.path.join(self.root, self.maskfolder , self.masks[idx])
img = Image.open(img_path).convert("RGB")
# The problem is that augmentation could remove all the visible labels here
img, mask = run_augmentations(img, mask)
mask = Image.open(mask_path)
mask = np.array(mask)
obj_ids = np.unique(mask)
obj_ids = obj_ids[1:]
masks = mask == obj_ids[:, None, None]
num_objs = len(obj_ids)
boxes = []
for i in range(num_objs):
pos = np.where(masks[i])
xmin = np.min(pos[1])
xmax = np.max(pos[1])
ymin = np.min(pos[0])
ymax = np.max(pos[0])
if xmin == xmax or ymin == ymax:
print("WARNING ZERO SIZE BOUNDING BOX AFTER AUGMENTATION")
print(img_path)
boxes.append([xmin, ymin, xmax, ymax])
if len(boxes) == 0:
print("WARNING ZERO BOUNDING BOXES AFTER AUGMENTATION")
print(img_path)
...
EDIT: My best attempt if that if there are no labels after augmentation return the image and mask without any augmentation.
