def __getitem__(self, index): ''' :param index: :return: image: tensor (H, W, channel=1) ''' def _read_image(path): cell = 8 input_image = cv2.imread(path) # print(f"path: {path}, image: {image}") # print(f"path: {path}, image: {input_image.shape}") input_image = cv2.resize(input_image, (self.sizer[1], self.sizer[0]), interpolation=cv2.INTER_AREA) H, W = input_image.shape[0], input_image.shape[1] # H = H//cell*cell # W = W//cell*cell # input_image = input_image[:H,:W,:] input_image = cv2.cvtColor(input_image, cv2.COLOR_RGB2GRAY) input_image = input_image.astype('float32') / 255.0 return input_image def _preprocess(image): if self.transforms is not None: image = self.transforms(image) return image def get_labels_gaussian(pnts, subpixel=False): heatmaps = np.zeros((H, W)) if subpixel: print("pnt: ", pnts.shape) for center in pnts: heatmaps = self.putGaussianMaps(center, heatmaps) else: aug_par = {'photometric': {}} aug_par['photometric']['enable'] = True aug_par['photometric']['params'] = self.config[ 'gaussian_label']['params'] augmentation = self.ImgAugTransform(**aug_par) # get label_2D labels = points_to_2D(pnts, H, W) labels = labels[:, :, np.newaxis] heatmaps = augmentation(labels) # warped_labels_gaussian = torch.tensor(heatmaps).float().view(-1, H, W) warped_labels_gaussian = torch.tensor(heatmaps).type( torch.FloatTensor).view(-1, H, W) warped_labels_gaussian[warped_labels_gaussian > 1.] = 1. return warped_labels_gaussian from datasets.data_tools import np_to_tensor # def np_to_tensor(img, H, W): # img = torch.tensor(img).type(torch.FloatTensor).view(-1, H, W) # return img from datasets.data_tools import warpLabels def imgPhotometric(img): """ :param img: numpy (H, W) :return: """ augmentation = self.ImgAugTransform(**self.config['augmentation']) img = img[:, :, np.newaxis] img = augmentation(img) cusAug = self.customizedTransform() img = cusAug(img, **self.config['augmentation']) return img def points_to_2D(pnts, H, W): labels = np.zeros((H, W)) pnts = pnts.astype(int) labels[pnts[:, 1], pnts[:, 0]] = 1 return labels to_floatTensor = lambda x: torch.tensor(x).type(torch.FloatTensor) from numpy.linalg import inv sample = self.samples[index] sample = self.format_sample(sample) input = {} input.update(sample) # image # img_o = _read_image(self.get_img_from_sample(sample)) img_o = _read_image(sample['image']) H, W = img_o.shape[0], img_o.shape[1] # print(f"image: {image.shape}") img_aug = img_o.copy() if (self.enable_photo_train == True and self.action == 'train') or (self.enable_photo_val and self.action == 'val'): img_aug = imgPhotometric(img_o) # numpy array (H, W, 1) # img_aug = _preprocess(img_aug[:,:,np.newaxis]) img_aug = torch.tensor(img_aug, dtype=torch.float32).view(-1, H, W) valid_mask = self.compute_valid_mask(torch.tensor([H, W]), inv_homography=torch.eye(3)) input.update({'image': img_aug}) input.update({'valid_mask': valid_mask}) if self.config['homography_adaptation']['enable']: # img_aug = torch.tensor(img_aug) homoAdapt_iter = self.config['homography_adaptation']['num'] homographies = np.stack([ self.sample_homography(np.array([2, 2]), shift=-1, **self.config['homography_adaptation'] ['homographies']['params']) for i in range(homoAdapt_iter) ]) ##### use inverse from the sample homography homographies = np.stack( [inv(homography) for homography in homographies]) homographies[0, :, :] = np.identity(3) # homographies_id = np.stack([homographies_id, homographies])[:-1,...] ###### homographies = torch.tensor(homographies, dtype=torch.float32) inv_homographies = torch.stack([ torch.inverse(homographies[i, :, :]) for i in range(homoAdapt_iter) ]) # images warped_img = self.inv_warp_image_batch( img_aug.squeeze().repeat(homoAdapt_iter, 1, 1, 1), inv_homographies, mode='bilinear').unsqueeze(0) warped_img = warped_img.squeeze() # masks valid_mask = self.compute_valid_mask( torch.tensor([H, W]), inv_homography=inv_homographies, erosion_radius=self.config['augmentation']['homographic'] ['valid_border_margin']) input.update({ 'image': warped_img, 'valid_mask': valid_mask, 'image_2D': img_aug }) input.update({ 'homographies': homographies, 'inv_homographies': inv_homographies }) # laebls if self.labels: pnts = np.load(sample['points'])['pts'] # pnts = pnts.astype(int) # labels = np.zeros_like(img_o) # labels[pnts[:, 1], pnts[:, 0]] = 1 labels = points_to_2D(pnts, H, W) labels_2D = to_floatTensor(labels[np.newaxis, :, :]) input.update({'labels_2D': labels_2D}) ## residual labels_res = torch.zeros((2, H, W)).type(torch.FloatTensor) input.update({'labels_res': labels_res}) if (self.enable_homo_train == True and self.action == 'train') or (self.enable_homo_val and self.action == 'val'): homography = self.sample_homography( np.array([2, 2]), shift=-1, **self.config['augmentation']['homographic']['params']) ##### use inverse from the sample homography homography = inv(homography) ###### inv_homography = inv(homography) inv_homography = torch.tensor(inv_homography).to(torch.float32) homography = torch.tensor(homography).to(torch.float32) # img = torch.from_numpy(img) warped_img = self.inv_warp_image(img_aug.squeeze(), inv_homography, mode='bilinear').unsqueeze(0) # warped_img = warped_img.squeeze().numpy() # warped_img = warped_img[:,:,np.newaxis] ##### check: add photometric ##### # labels = torch.from_numpy(labels) # warped_labels = self.inv_warp_image(labels.squeeze(), inv_homography, mode='nearest').unsqueeze(0) ##### check ##### warped_set = warpLabels(pnts, H, W, homography) warped_labels = warped_set['labels'] # if self.transform is not None: # warped_img = self.transform(warped_img) valid_mask = self.compute_valid_mask( torch.tensor([H, W]), inv_homography=inv_homography, erosion_radius=self.config['augmentation']['homographic'] ['valid_border_margin']) input.update({ 'image': warped_img, 'labels_2D': warped_labels, 'valid_mask': valid_mask }) if self.config['warped_pair']['enable']: homography = self.sample_homography( np.array([2, 2]), shift=-1, **self.config['warped_pair']['params']) ##### use inverse from the sample homography homography = np.linalg.inv(homography) ##### inv_homography = np.linalg.inv(homography) homography = torch.tensor(homography).type(torch.FloatTensor) inv_homography = torch.tensor(inv_homography).type( torch.FloatTensor) # photometric augmentation from original image # warp original image warped_img = torch.tensor(img_o, dtype=torch.float32) warped_img = self.inv_warp_image(warped_img.squeeze(), inv_homography, mode='bilinear').unsqueeze(0) if (self.enable_photo_train == True and self.action == 'train') or (self.enable_photo_val and self.action == 'val'): warped_img = imgPhotometric( warped_img.numpy().squeeze()) # numpy array (H, W, 1) warped_img = torch.tensor(warped_img, dtype=torch.float32) pass warped_img = warped_img.view(-1, H, W) # warped_labels = warpLabels(pnts, H, W, homography) warped_set = warpLabels(pnts, H, W, homography, bilinear=True) warped_labels = warped_set['labels'] warped_res = warped_set['res'] warped_res = warped_res.transpose(1, 2).transpose(0, 1) # print("warped_res: ", warped_res.shape) if self.gaussian_label: # print("do gaussian labels!") # warped_labels_gaussian = get_labels_gaussian(warped_set['warped_pnts'].numpy()) from utils.var_dim import squeezeToNumpy # warped_labels_bi = self.inv_warp_image(labels_2D.squeeze(), inv_homography, mode='nearest').unsqueeze(0) # bilinear, nearest warped_labels_bi = warped_set['labels_bi'] warped_labels_gaussian = self.gaussian_blur( squeezeToNumpy(warped_labels_bi)) warped_labels_gaussian = np_to_tensor( warped_labels_gaussian, H, W) input['warped_labels_gaussian'] = warped_labels_gaussian input.update({'warped_labels_bi': warped_labels_bi}) input.update({ 'warped_img': warped_img, 'warped_labels': warped_labels, 'warped_res': warped_res }) # print('erosion_radius', self.config['warped_pair']['valid_border_margin']) valid_mask = self.compute_valid_mask( torch.tensor([H, W]), inv_homography=inv_homography, erosion_radius=self.config['warped_pair'] ['valid_border_margin']) # can set to other value input.update({'warped_valid_mask': valid_mask}) input.update({ 'homographies': homography, 'inv_homographies': inv_homography }) # labels = self.labels2Dto3D(self.cell_size, labels) # labels = torch.from_numpy(labels[np.newaxis,:,:]) # input.update({'labels': labels}) if self.gaussian_label: # warped_labels_gaussian = get_labels_gaussian(pnts) labels_gaussian = self.gaussian_blur(squeezeToNumpy(labels_2D)) labels_gaussian = np_to_tensor(labels_gaussian, H, W) input['labels_2D_gaussian'] = labels_gaussian name = sample['name'] to_numpy = False if to_numpy: image = np.array(img) input.update({'name': name, 'scene_name': "./"}) # dummy scene name return input
def __getitem__(self, index): """ :param index: :return: labels_2D: tensor(1, H, W) image: tensor(1, H, W) """ def checkSat(img, name=""): if img.max() > 1: print(name, img.max()) elif img.min() < 0: print(name, img.min()) def imgPhotometric(img): """ :param img: numpy (H, W) :return: """ augmentation = self.ImgAugTransform(**self.config["augmentation"]) img = img[:, :, np.newaxis] img = augmentation(img) cusAug = self.customizedTransform() img = cusAug(img, **self.config["augmentation"]) return img def get_labels(pnts, H, W): labels = torch.zeros(H, W) # print('--2', pnts, pnts.size()) # pnts_int = torch.min(pnts.round().long(), torch.tensor([[H-1, W-1]]).long()) pnts_int = torch.min(pnts.round().long(), torch.tensor([[W - 1, H - 1]]).long()) # print('--3', pnts_int, pnts_int.size()) labels[pnts_int[:, 1], pnts_int[:, 0]] = 1 return labels def get_label_res(H, W, pnts): quan = lambda x: x.round().long() labels_res = torch.zeros(H, W, 2) # pnts_int = torch.min(pnts.round().long(), torch.tensor([[H-1, W-1]]).long()) labels_res[quan(pnts)[:, 1], quan(pnts)[:, 0], :] = pnts - pnts.round() # print("pnts max: ", quan(pnts).max(dim=0)) # print("labels_res: ", labels_res.shape) labels_res = labels_res.transpose(1, 2).transpose(0, 1) return labels_res from datasets.data_tools import np_to_tensor from utils.utils import filter_points from utils.var_dim import squeezeToNumpy sample = self.samples[index] img = load_as_float(sample["image"]) H, W = img.shape[0], img.shape[1] self.H = H self.W = W pnts = np.load(sample["points"]) # (y, x) pnts = torch.tensor(pnts).float() pnts = torch.stack((pnts[:, 1], pnts[:, 0]), dim=1) # (x, y) pnts = filter_points(pnts, torch.tensor([W, H])) sample = {} # print('pnts: ', pnts[:5]) # print('--1', pnts) labels_2D = get_labels(pnts, H, W) sample.update({"labels_2D": labels_2D.unsqueeze(0)}) # assert Hc == round(Hc) and Wc == round(Wc), "Input image size not fit in the block size" if (self.config["augmentation"]["photometric"]["enable_train"] and self.action == "training") or ( self.config["augmentation"]["photometric"]["enable_val"] and self.action == "validation"): # print('>>> Photometric aug enabled for %s.'%self.action) # augmentation = self.ImgAugTransform(**self.config["augmentation"]) img = imgPhotometric(img) else: # print('>>> Photometric aug disabled for %s.'%self.action) pass if not ((self.config["augmentation"]["homographic"]["enable_train"] and self.action == "training") or (self.config["augmentation"]["homographic"]["enable_val"] and self.action == "validation")): # print('<<< Homograpy aug disabled for %s.'%self.action) img = img[:, :, np.newaxis] # labels = labels.view(-1,H,W) if self.transform is not None: img = self.transform(img) sample["image"] = img # sample = {'image': img, 'labels_2D': labels} valid_mask = self.compute_valid_mask(torch.tensor([H, W]), inv_homography=torch.eye(3)) sample.update({"valid_mask": valid_mask}) labels_res = get_label_res(H, W, pnts) pnts_post = pnts # pnts_for_gaussian = pnts else: # print('>>> Homograpy aug enabled for %s.'%self.action) # img_warp = img from utils.utils import homography_scaling_torch as homography_scaling from numpy.linalg import inv homography = self.sample_homography( np.array([2, 2]), shift=-1, **self.config["augmentation"]["homographic"]["params"], ) ##### use inverse from the sample homography homography = inv(homography) ###### homography = torch.tensor(homography).float() inv_homography = homography.inverse() img = torch.from_numpy(img) warped_img = self.inv_warp_image(img.squeeze(), inv_homography, mode="bilinear") warped_img = warped_img.squeeze().numpy() warped_img = warped_img[:, :, np.newaxis] # labels = torch.from_numpy(labels) # warped_labels = self.inv_warp_image(labels.squeeze(), inv_homography, mode='nearest').unsqueeze(0) warped_pnts = self.warp_points( pnts, homography_scaling(homography, H, W)) warped_pnts = filter_points(warped_pnts, torch.tensor([W, H])) # pnts = warped_pnts[:, [1, 0]] # pnts_for_gaussian = warped_pnts # warped_labels = torch.zeros(H, W) # warped_labels[warped_pnts[:, 1], warped_pnts[:, 0]] = 1 # warped_labels = warped_labels.view(-1, H, W) if self.transform is not None: warped_img = self.transform(warped_img) # sample = {'image': warped_img, 'labels_2D': warped_labels} sample["image"] = warped_img valid_mask = self.compute_valid_mask( torch.tensor([H, W]), inv_homography=inv_homography, erosion_radius=self.config["augmentation"]["homographic"] ["valid_border_margin"], ) # can set to other value sample.update({"valid_mask": valid_mask}) labels_2D = get_labels(warped_pnts, H, W) sample.update({"labels_2D": labels_2D.unsqueeze(0)}) labels_res = get_label_res(H, W, warped_pnts) pnts_post = warped_pnts if self.gaussian_label: # warped_labels_gaussian = get_labels_gaussian(pnts) from datasets.data_tools import get_labels_bi labels_2D_bi = get_labels_bi(pnts_post, H, W) labels_gaussian = self.gaussian_blur(squeezeToNumpy(labels_2D_bi)) labels_gaussian = np_to_tensor(labels_gaussian, H, W) sample["labels_2D_gaussian"] = labels_gaussian # add residua sample.update({"labels_res": labels_res}) ### code for warped image if self.config["warped_pair"]["enable"]: from datasets.data_tools import warpLabels homography = self.sample_homography( np.array([2, 2]), shift=-1, **self.config["warped_pair"]["params"]) ##### use inverse from the sample homography homography = np.linalg.inv(homography) ##### inv_homography = np.linalg.inv(homography) homography = torch.tensor(homography).type(torch.FloatTensor) inv_homography = torch.tensor(inv_homography).type( torch.FloatTensor) # photometric augmentation from original image # warp original image warped_img = img.type(torch.FloatTensor) warped_img = self.inv_warp_image(warped_img.squeeze(), inv_homography, mode="bilinear").unsqueeze(0) if (self.enable_photo_train == True and self.action == "train") or (self.enable_photo_val and self.action == "val"): warped_img = imgPhotometric( warped_img.numpy().squeeze()) # numpy array (H, W, 1) warped_img = torch.tensor(warped_img, dtype=torch.float32) pass warped_img = warped_img.view(-1, H, W) # warped_labels = warpLabels(pnts, H, W, homography) warped_set = warpLabels(pnts, H, W, homography, bilinear=True) warped_labels = warped_set["labels"] warped_res = warped_set["res"] warped_res = warped_res.transpose(1, 2).transpose(0, 1) # print("warped_res: ", warped_res.shape) if self.gaussian_label: # print("do gaussian labels!") # warped_labels_gaussian = get_labels_gaussian(warped_set['warped_pnts'].numpy()) # warped_labels_bi = self.inv_warp_image(labels_2D.squeeze(), inv_homography, mode='nearest').unsqueeze(0) # bilinear, nearest warped_labels_bi = warped_set["labels_bi"] warped_labels_gaussian = self.gaussian_blur( squeezeToNumpy(warped_labels_bi)) warped_labels_gaussian = np_to_tensor(warped_labels_gaussian, H, W) sample["warped_labels_gaussian"] = warped_labels_gaussian sample.update({"warped_labels_bi": warped_labels_bi}) sample.update({ "warped_img": warped_img, "warped_labels": warped_labels, "warped_res": warped_res, }) # print('erosion_radius', self.config['warped_pair']['valid_border_margin']) valid_mask = self.compute_valid_mask( torch.tensor([H, W]), inv_homography=inv_homography, erosion_radius=self.config["warped_pair"] ["valid_border_margin"], ) # can set to other value sample.update({"warped_valid_mask": valid_mask}) sample.update({ "homographies": homography, "inv_homographies": inv_homography }) # labels = self.labels2Dto3D(self.cell_size, labels) # labels = torch.from_numpy(labels[np.newaxis,:,:]) # input.update({'labels': labels}) ### code for warped image # if self.config['gaussian_label']['enable']: # heatmaps = np.zeros((H, W)) # # for center in pnts_int.numpy(): # for center in pnts[:, [1, 0]].numpy(): # # print("put points: ", center) # heatmaps = self.putGaussianMaps(center, heatmaps) # # import matplotlib.pyplot as plt # # plt.figure(figsize=(5, 10)) # # plt.subplot(211) # # plt.imshow(heatmaps) # # plt.colorbar() # # plt.subplot(212) # # plt.imshow(np.squeeze(warped_labels.numpy())) # # plt.show() # # import time # # time.sleep(500) # # results = self.pool.map(self.putGaussianMaps_par, warped_pnts.numpy()) # warped_labels_gaussian = torch.from_numpy(heatmaps).view(-1, H, W) # warped_labels_gaussian[warped_labels_gaussian>1.] = 1. # sample['labels_2D_gaussian'] = warped_labels_gaussian if self.getPts: sample.update({"pts": pnts}) return sample