import numpy as np from torch.utils.data import Dataset from scipy.ndimage import gaussian_filter import cv2 mean = [0.485, 0.456, 0.406] std = [0.229, 0.224, 0.225] class CornersDataset(Dataset): def __init__(self, image_size=256, inference=False): super(CornersDataset, self).__init__() self.image_size = image_size self.inference = inference self._data_names = [] def __len__(self): raise len(self._data_names) def __getitem__(self, idx): raise NotImplementedError def process_data(self, data): img = data['image'] corners = data['corners'] annot = data['annot'] # pre-process the image to use ImageNet-pretrained backbones img = img.transpose((2, 0, 1)) raw_img = img.copy() img = (img - np.array(mean)[:, np.newaxis, np.newaxis]) / np.array(std)[:, np.newaxis, np.newaxis] img = img.astype(np.float32) corners = np.array(corners) all_data = { "annot": annot, "name": data['name'], 'img': img, 'annot_path': data['annot_path'], 'img_path': data['img_path'], 'det_path': data['det_path'], 'raw_img': raw_img, } # corner labels for training if not self.inference: pixel_labels, gauss_labels = self.get_corner_labels(corners) all_data['pixel_labels'] = pixel_labels all_data['gauss_labels'] = gauss_labels return all_data def get_corner_labels(self, corners): labels = np.zeros((self.image_size, self.image_size)) corners = corners.round() xint, yint = corners[:, 0].astype(np.int), corners[:, 1].astype(np.int) labels[yint, xint] = 1 gauss_labels = gaussian_filter(labels, sigma=2) gauss_labels = gauss_labels / gauss_labels.max() return labels, gauss_labels def resize_data(self, image, annot, det_corners): new_image = cv2.resize(image, (self.image_size, self.image_size)) new_annot = {} r = self.image_size / 256 for c, connections in annot.items(): new_c = tuple(np.array(c) * r) new_connections = [other_c * r for other_c in connections] new_annot[new_c] = new_connections new_dets = det_corners * r return new_image, new_annot, new_dets def random_aug_annot(self, img, annot, det_corners=None): # do random flipping img, annot, det_corners = self.random_flip(img, annot, det_corners) # prepare random augmentation parameters (only do random rotation for now) theta = np.random.randint(0, 360) / 360 * np.pi * 2 r = self.image_size / 256 origin = [127 * r, 127 * r] p1_new = [127 * r + 100 * np.sin(theta) * r, 127 * r - 100 * np.cos(theta) * r] p2_new = [127 * r + 100 * np.cos(theta) * r, 127 * r + 100 * np.sin(theta) * r] p1_old = [127 * r, 127 * r - 100 * r] # y_axis p2_old = [127 * r + 100 * r, 127 * r] # x_axis pts1 = np.array([origin, p1_old, p2_old]).astype(np.float32) pts2 = np.array([origin, p1_new, p2_new]).astype(np.float32) M_rot = cv2.getAffineTransform(pts1, pts2) # Combine annotation corners and detection corners all_corners = list(annot.keys()) if det_corners is not None: for i in range(det_corners.shape[0]): all_corners.append(tuple(det_corners[i])) all_corners_ = np.array(all_corners) # Do the corner transform within a big matrix transformation corner_mapping = dict() ones = np.ones([all_corners_.shape[0], 1]) all_corners_ = np.concatenate([all_corners_, ones], axis=-1) aug_corners = np.matmul(M_rot, all_corners_.T).T for idx, corner in enumerate(all_corners): corner_mapping[corner] = aug_corners[idx] # If the transformed geometry goes beyond image boundary, we simply re-do the augmentation new_corners = np.array(list(corner_mapping.values())) if new_corners.min() <= 0 or new_corners.max() >= (self.image_size - 1): # return self.random_aug_annot(img, annot, det_corners) return img, annot, None, det_corners # build the new annot dict aug_annot = dict() for corner, connections in annot.items(): new_corner = corner_mapping[corner] tuple_new_corner = tuple(new_corner) aug_annot[tuple_new_corner] = list() for to_corner in connections: aug_annot[tuple_new_corner].append(corner_mapping[tuple(to_corner)]) # Also transform the image correspondingly rows, cols, ch = img.shape new_img = cv2.warpAffine(img, M_rot, (cols, rows), borderValue=(255, 255, 255)) y_start = (new_img.shape[0] - self.image_size) // 2 x_start = (new_img.shape[1] - self.image_size) // 2 aug_img = new_img[y_start:y_start + self.image_size, x_start:x_start + self.image_size, :] if det_corners is None: return aug_img, aug_annot, corner_mapping, None else: aug_det_corners = list() for corner in det_corners: new_corner = corner_mapping[tuple(corner)] aug_det_corners.append(new_corner) aug_det_corners = np.array(aug_det_corners) return aug_img, aug_annot, corner_mapping, aug_det_corners def random_flip(self, img, annot, det_corners): height, width, _ = img.shape rand_int = np.random.randint(0, 4) if rand_int == 0: return img, annot, det_corners all_corners = list(annot.keys()) if det_corners is not None: for i in range(det_corners.shape[0]): all_corners.append(tuple(det_corners[i])) new_corners = np.array(all_corners) if rand_int == 1: img = img[:, ::-1, :] new_corners[:, 0] = width - new_corners[:, 0] elif rand_int == 2: img = img[::-1, :, :] new_corners[:, 1] = height - new_corners[:, 1] else: img = img[::-1, ::-1, :] new_corners[:, 0] = width - new_corners[:, 0] new_corners[:, 1] = height - new_corners[:, 1] new_corners = np.clip(new_corners, 0, self.image_size - 1) # clip into [0, 255] corner_mapping = dict() for idx, corner in enumerate(all_corners): corner_mapping[corner] = new_corners[idx] aug_annot = dict() for corner, connections in annot.items(): new_corner = corner_mapping[corner] tuple_new_corner = tuple(new_corner) aug_annot[tuple_new_corner] = list() for to_corner in connections: aug_annot[tuple_new_corner].append(corner_mapping[tuple(to_corner)]) if det_corners is not None: aug_det_corners = list() for corner in det_corners: new_corner = corner_mapping[tuple(corner)] aug_det_corners.append(new_corner) det_corners = np.array(aug_det_corners) return img, aug_annot, det_corners