deep3d-pytorch/dataloader.py at master · anishmadan23/deep3d-pytorch · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
import os
import torch
import torchvision

import numpy as np
import skimage.io as io
import torch.utils.data as data

from skimage.transform import resize

class MyDataset(data.Dataset):
	def __init__(self, root, in_transforms = None, orig_size = (384, 1280),small_size=(96,320)):
		self.leftpath = os.path.join(root, 'left')
		self.leftimg = os.listdir(self.leftpath)

		self.rightpath = os.path.join(root, 'right')
		self.rightimg = os.listdir(self.rightpath)

		self.leftimg.sort()
		self.rightimg.sort()

		self.orig_size = orig_size
		self.small_size= small_size

	def __len__(self):
		return len(self.leftimg)

	def __getitem__(self, index):
		leftImage = io.imread(os.path.join(self.leftpath, self.leftimg[index]))
		# print(leftImage.shape)
		leftImage_orig = resize(leftImage, self.orig_size) / 255.0

		leftImage_small = resize(leftImage, self.small_size) / 255.0

		rightImage_orig = io.imread(os.path.join(self.rightpath, self.rightimg[index]))
		rightImage_orig = resize(rightImage_orig, self.orig_size) /255.0

		left_orig = torch.from_numpy(leftImage_orig)
		left_orig = left_orig.permute([-1,0,1])

		left_small = torch.from_numpy(leftImage_small)
		left_small = left_small.permute([-1,0,1])

		right_orig = torch.from_numpy(rightImage_orig)
		right_orig = right_orig.permute([-1,0,1])

		return left_orig, left_small, right_orig


# data_obj = MyDataset('./data/train/')
# train_dataloader = data.DataLoader(data_obj, batch_size = 4, shuffle = True)