def _apply_mosaic(self, images, targets):
assert len(images) == 4 and len(targets) == 4
sw, sh = self._img_size
c = images[0].shape[2]
sum_rgb = np.zeros([images[0].ndim, ])
for img in images:
sum_rgb += np.array(cv2.mean(img))[:3]
mean_rgb = sum_rgb / len(images)
img4 = np.full((sh * 2, sw * 2, c), mean_rgb, dtype=np.uint8) # base image with 4 tiles
offsets = [(0, 0), (sw, 0), (0, sh), (sw, sh)]
target4 = ParamList((sw, sh))
for i, img, target in zip(range(4), images, targets):
h, w, _ = img.shape
pad_w = int(sw - w) // 2
pad_h = int(sh - h) // 2
y_st = pad_h + offsets[i][1]
x_st = pad_w + offsets[i][0]
img4[y_st:y_st + h, x_st:x_st + w] = img
bbox = target.get_field('bbox')
bbox[:, 0::2] += x_st
bbox[:, 1::2] += y_st
target.update_field('bbox', bbox)
# np.clip(bbox[:, 0::2], 0, 2 * sw, out=bbox[:, 0::2]) # use with random_affine
# np.clip(bbox[:, 1::2], 0, 2 * sh, out=bbox[:, 1::2])
target4.merge(target)
raff = transforms.RandomAffine2D()
param = {
'border': (-sh//2, -sw//2)
}
param.update(self._config)
return raff(img4, target4, **param)
def collate_fn(batch):
img, target, path, shape = zip(*batch) # transposed
ntarget = ParamList((None, None))
for i, t in enumerate(target):
id = t.get_field('img_id')
id[:, ] = i
t.update_field('img_id', id)
ntarget.merge(t)
# ntarget.to_tensor()
return torch.stack(img, 0), ntarget, path, shape