def __getitem__(self, index): img_idx = self.lists[index] img = self.imgs[img_idx].transpose(2, 1, 0) #img = self.imgs[img_idx] dpt = self.dpts[img_idx].transpose(1, 0) #dpt = self.dpts[img_idx] #image = Image.fromarray(np.uint8(img)) #depth = Image.fromarray(np.uint8(dpt)) #image.save('img1.png') #input_transform = transforms.Compose([flow_transforms.Scale(228)]) #input_transform = transforms.Compose([flow_transforms.ArrayToTensor()]) input_transform = transforms.Compose( [flow_transforms.Scale(228), flow_transforms.ArrayToTensor()]) #target_depth_transform = transforms.Compose([flow_transforms.Scale(228)]) #target_depth_transform = transforms.Compose([flow_transforms.ArrayToTensor()]) target_depth_transform = transforms.Compose([ flow_transforms.Scale_Single(228), flow_transforms.ArrayToTensor() ]) img = input_transform(img) dpt = target_depth_transform(dpt) #image = Image.fromarray(np.uint8(img)) #image.save('img2.png') return img, dpt
def load_data(is_training = True): global IMAGE_NUM train_idx_path = "data/trainNdxs.txt" test_idx_path = "data/testNdxs.txt" input_rgb_images_dir = 'data/nyu_datasets_changed/input/' target_depth_images_dir = 'data/nyu_datasets_changed/target_depths/' target_labels_images_dir = 'data/nyu_datasets_changed/labels_38/' data_path = "data/nyu_depth_v2_labeled.mat" train_idx = np.loadtxt(train_idx_path, dtype = 'int') test_idx = np.loadtxt(test_idx_path, dtype = 'int') input_transform = flow_transforms.Compose([flow_transforms.Scale(120)]) target_depth_transform = flow_transforms.Compose([flow_transforms.Scale_Single(60)]) target_labels_transform = flow_transforms.Compose([]) co_transform=flow_transforms.Compose([ flow_transforms.RandomRotate(4), flow_transforms.RandomCrop((480,640)), flow_transforms.RandomVerticalFlip() ]) data = [] if is_training: data = ListDataset(data_path,train_idx,input_transform,target_depth_transform, target_labels_transform, co_transform) else: data = ListDataset(data_path, test_idx, input_transform, target_depth_transform, target_labels_transform) IMAGE_NUM = len(data) return data
listing = random.sample(os.listdir(input_rgb_images_dir), 1449) train_listing = listing[:min(NUM_TRAIN, train_on)] val_listing = listing[NUM_TRAIN:min(NUM_VAL + NUM_TRAIN, val_on + NUM_TRAIN)] test_listing = listing[NUM_VAL + NUM_TRAIN:min(NUM_VAL + NUM_TRAIN + NUM_TEST, test_on + NUM_VAL + NUM_TRAIN)] data_dir = (input_rgb_images_dir, target_depth_images_dir, target_labels_images_dir) input_transform = transforms.Compose( [flow_transforms.Scale(228), flow_transforms.ArrayToTensor()]) target_depth_transform = transforms.Compose( [flow_transforms.Scale_Single(228), flow_transforms.ArrayToTensor()]) target_labels_transform = transforms.Compose([flow_transforms.ArrayToTensor()]) ##Apply this transform on input, ground truth depth images and labeled images co_transform = flow_transforms.Compose([ flow_transforms.RandomCrop((480, 640)), flow_transforms.RandomHorizontalFlip() ]) ##Splitting in train, val and test sets [No data augmentation on val and test, only on train] train_dataset = ListDataset(data_dir,train_listing,input_transform,target_depth_transform,\ target_labels_transform,co_transform)