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)