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