import os
import torch
import torch.utils.data as data
import torch
import torchvision.transforms as transforms
import random
from PIL import Image, ImageOps
import numpy as np
import preprocess
IMG_EXTENSIONS = [
'.jpg', '.JPG', '.jpeg', '.JPEG',
'.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
]
def is_image_file(filename):
return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
def default_loader(path):
return Image.open(path).convert('RGB')
def disparity_loader(path):
return Image.open(path)
class myImageFloder(data.Dataset):
def __init__(self, left, right, left_disparity, training, loader=default_loader, dploader= disparity_loader):
self.left = left
self.right = right
self.disp_L = left_disparity
self.loader = loader
self.dploader = dploader
self.training = training
def __getitem__(self, index):
left = self.left[index]
right = self.right[index]
disp_L= self.disp_L[index]
left_img = self.loader(left)
right_img = self.loader(right)
dataL = self.dploader(disp_L)
if self.training:
w, h = left_img.size
th, tw = 256, 512
x1 = random.randint(0, w - tw)
y1 = random.randint(0, h - th)
left_img = left_img.crop((x1, y1, x1 + tw, y1 + th))
right_img = right_img.crop((x1, y1, x1 + tw, y1 + th))
dataL = np.ascontiguousarray(dataL,dtype=np.float32)/256
dataL = dataL[y1:y1 + th, x1:x1 + tw]
processed = preprocess.get_transform(augment=False)
left_img = processed(left_img)
right_img = processed(right_img)
return left_img, right_img, dataL
else:
w, h = left_img.size
left_img = left_img.crop((w-1232, h-368, w, h))
right_img = right_img.crop((w-1232, h-368, w, h))
w1, h1 = left_img.size
dataL = dataL.crop((w-1232, h-368, w, h))
dataL = np.ascontiguousarray(dataL,dtype=np.float32)/256
processed = preprocess.get_transform(augment=False)
left_img = processed(left_img)
right_img = processed(right_img)
return left_img, right_img, dataL
def __len__(self):
return len(self.left)
