def __getitem__(self, index):
left_img = self.load_image(
os.path.join(self.datapath, self.left_filenames[index]))
right_img = self.load_image(
os.path.join(self.datapath, self.right_filenames[index]))
disparity = self.load_disp(
os.path.join(self.datapath, self.disp_filenames[index]))
if self.training:
w, h = left_img.size
crop_w, crop_h = 512, 256
x1 = random.randint(0, w - crop_w)
y1 = random.randint(0, h - crop_h)
# random crop
left_img = left_img.crop((x1, y1, x1 + crop_w, y1 + crop_h))
right_img = right_img.crop((x1, y1, x1 + crop_w, y1 + crop_h))
disparity = disparity[y1:y1 + crop_h, x1:x1 + crop_w]
# to tensor, normalize
processed = get_transform()
left_img = processed(left_img)
right_img = processed(right_img)
return {"left": left_img, "right": right_img, "disp": disparity}
else:
w, h = left_img.size
crop_w, crop_h = 512, 256
left_img = left_img.crop((w - crop_w, h - crop_h, w, h))
right_img = right_img.crop((w - crop_w, h - crop_h, w, h))
disparity = disparity[h - crop_h:h, w - crop_w:w]
processed = get_transform()
left_img = processed(left_img)
right_img = processed(right_img)
return {
"left": left_img,
"right": right_img,
"disp": disparity,
"top_pad": 0,
"right_pad": 0
}
import numpy as np
from PIL import Image
from datasets.data_io import get_transform
from models.bgnet import BGNet
from models.bgnet_plus import BGNet_Plus
model = BGNet_Plus().cuda()
checkpoint = torch.load('models/Sceneflow-IRS-BGNet-Plus.pth',
map_location=lambda storage, loc: storage)
model.load_state_dict(checkpoint)
model.eval()
left_img = Image.open('sample/im0.png').convert('L')
right_img = Image.open('sample/im1.png').convert('L')
w, h = left_img.size
h1 = h % 64
w1 = w % 64
h1 = h - h1
w1 = w - w1
h1 = int(h1)
w1 = int(w1)
left_img = left_img.resize((w1, h1), Image.ANTIALIAS)
right_img = right_img.resize((w1, h1), Image.ANTIALIAS)
left_img = np.ascontiguousarray(left_img, dtype=np.float32)
right_img = np.ascontiguousarray(right_img, dtype=np.float32)
preprocess = get_transform()
left_img = preprocess(left_img)
right_img = preprocess(right_img)
pred, _ = model(left_img.unsqueeze(0).cuda(), right_img.unsqueeze(0).cuda())
pred = pred[0].data.cpu().numpy() * 256
skimage.io.imsave('sample_disp.png', pred.astype('uint16'))
def __getitem__(self, index):
left_img = self.load_image(
os.path.join(self.datapath, self.left_filenames[index]))
right_img = self.load_image(
os.path.join(self.datapath, self.right_filenames[index]))
if self.disp_filenames: # has disparity ground truth
disparity = self.load_disp(
os.path.join(self.datapath, self.disp_filenames[index]))
else:
disparity = None
if self.training:
w, h = left_img.size
crop_w, crop_h = self.crop_width, self.crop_height
x1 = random.randint(0, w - crop_w)
y1 = random.randint(0, h - crop_h)
# random crop
left_img = left_img.crop((x1, y1, x1 + crop_w, y1 + crop_h))
right_img = right_img.crop((x1, y1, x1 + crop_w, y1 + crop_h))
disparity = disparity[y1:y1 + crop_h, x1:x1 + crop_w]
# to tensor, normalize
processed = get_transform()
left_img = processed(left_img)
right_img = processed(right_img)
return {
"left": left_img,
"right": right_img,
"disparity": disparity
}
else:
w, h = left_img.size
# normalize
processed = get_transform()
left_img = processed(left_img).numpy()
right_img = processed(right_img).numpy()
# pad to size 1248x384
top_pad = self.test_crop_height - h
right_pad = self.test_crop_width - w
assert top_pad > 0 and right_pad > 0
# pad images
left_img = np.lib.pad(left_img,
((0, 0), (top_pad, 0), (0, right_pad)),
mode='constant',
constant_values=0)
right_img = np.lib.pad(right_img,
((0, 0), (top_pad, 0), (0, right_pad)),
mode='constant',
constant_values=0)
# pad disparity gt
if disparity is not None:
assert len(disparity.shape) == 2
disparity = np.lib.pad(disparity,
((top_pad, 0), (0, right_pad)),
mode='constant',
constant_values=0)
if disparity is not None:
return {
"left": left_img,
"right": right_img,
"disparity": disparity,
"top_pad": top_pad,
"right_pad": right_pad
}
else:
return {
"left": left_img,
"right": right_img,
"top_pad": top_pad,
"right_pad": right_pad,
"left_filename": self.left_filenames[index],
"right_filename": self.right_filenames[index]
}
def __getitem__(self, index):
left_img = self.load_image(
os.path.join(self.datapath, self.left_filenames[index]))
right_img = self.load_image(
os.path.join(self.datapath, self.right_filenames[index]))
left_disp = self.load_disp(
os.path.join(self.datapath, self.left_disp_filenames[index]))
right_disp = self.load_disp(
os.path.join(self.datapath, self.right_disp_filenames[index]))
if self.training:
w, h = left_img.size
crop_w, crop_h = 512, 256
x1 = random.randint(0, w - crop_w)
y1 = random.randint(0, h - crop_h)
# random crop
left_img = left_img.crop((x1, y1, x1 + crop_w, y1 + crop_h))
right_img = right_img.crop((x1, y1, x1 + crop_w, y1 + crop_h))
left_disp = left_disp[y1:y1 + crop_h, x1:x1 + crop_w]
right_disp = right_disp[y1:y1 + crop_h, x1:x1 + crop_w]
# to tensor, normalize
processed = get_transform()
left_img = processed(left_img)
right_img = processed(right_img)
# augumentation
# if random.random()<0.5:
# left_img = torch.flip(left_img, [1])
# right_img = torch.flip(right_img, [1])
# left_disp = np.ascontiguousarray(np.flip(left_disp, 0))
# right_disp = np.ascontiguousarray(np.flip(right_disp, 0))
return {
"left": left_img,
"right": right_img,
"left_disp": left_disp,
"right_disp": right_disp
}
else:
w, h = left_img.size
crop_w, crop_h = 960, 512
left_img = left_img.crop((w - crop_w, h - crop_h, w, h))
right_img = right_img.crop((w - crop_w, h - crop_h, w, h))
disparity = left_disp[h - crop_h:h, w - crop_w:w]
disparity_right = right_disp[h - crop_h:h, w - crop_w:w]
processed = get_transform()
left_img = processed(left_img)
right_img = processed(right_img)
return {
"left": left_img,
"right": right_img,
"left_disp": disparity,
"right_disp": disparity_right,
"top_pad": 0,
"right_pad": 0
}
def __getitem__(self, index):
left_img = self.load_image(
os.path.join(self.datapath, self.left_filenames[index]))
right_img = self.load_image(
os.path.join(self.datapath, self.right_filenames[index]))
disparity = self.load_disp(
os.path.join(self.datapath, self.disp_filenames[index]))
# occ
occ = np.array(Image.open(self.occ_data[index])) != 255
disparity[occ] = 0.0
if self.training:
w, h = left_img.size
crop_w, crop_h = 512, 256
x1 = random.randint(0, w - crop_w)
y1 = random.randint(0, h - crop_h)
# random crop
left_img = left_img.crop((x1, y1, x1 + crop_w, y1 + crop_h))
right_img = right_img.crop((x1, y1, x1 + crop_w, y1 + crop_h))
disparity = disparity[y1:y1 + crop_h, x1:x1 + crop_w]
# to tensor, normalize
processed = get_transform()
left_img = processed(left_img)
right_img = processed(right_img)
return {
"left": left_img,
"right": right_img,
"disparity": disparity
}
else:
w, h = left_img.size
top_pad = math.ceil(h / 32) * 32 - h
right_pad = math.ceil(w / 32) * 32 - w
processed = get_transform()
left_img = processed(left_img)
right_img = processed(right_img)
# pad images
left_img = np.lib.pad(left_img,
((0, 0), (top_pad, 0), (0, right_pad)),
mode='constant',
constant_values=0)
right_img = np.lib.pad(right_img,
((0, 0), (top_pad, 0), (0, right_pad)),
mode='constant',
constant_values=0)
# pad disparity gt
if disparity is not None:
assert len(disparity.shape) == 2
disparity = np.lib.pad(disparity,
((top_pad, 0), (0, right_pad)),
mode='constant',
constant_values=0)
return {
"left": left_img,
"right": right_img,
"disparity": disparity,
"top_pad": top_pad,
"right_pad": right_pad
}
def __getitem__(self, index):
left_img = self.load_image(os.path.join(self.datapath, self.left_filenames[index]))
right_img = self.load_image(os.path.join(self.datapath, self.right_filenames[index]))
if self.disp_occ_filenames: # has disparity ground truth
disp_occ = self.load_disp(os.path.join(self.datapath, self.disp_occ_filenames[index]))
disp_noc = self.load_disp(os.path.join(self.datapath, self.disp_noc_filenames[index]))
else:
disp_occ = None
if self.training:
w, h = left_img.size
crop_w, crop_h = 512, 256
x1 = random.randint(0, w - crop_w)
y1 = random.randint(0, h - crop_h)
# random crop
left_img = left_img.crop((x1, y1, x1 + crop_w, y1 + crop_h))
right_img = right_img.crop((x1, y1, x1 + crop_w, y1 + crop_h))
disp_occ = disp_occ[y1:y1 + crop_h, x1:x1 + crop_w]
disp_noc = disp_noc[y1:y1 + crop_h, x1:x1 + crop_w]
occ_mask = ((disp_occ - disp_noc) > 0).astype(np.float32)
# to tensor, normalize
processed = get_transform()
left_img = processed(left_img)
right_img = processed(right_img)
# # augumentation
# if random.random() < 0.5:
# left_img = torch.flip(left_img, [1])
# right_img = torch.flip(right_img, [1])
# disp_occ = np.ascontiguousarray(np.flip(disp_occ, 0))
return {"left": left_img,
"right": right_img,
"left_disp": disp_occ,
"occ_mask": occ_mask}
else:
w, h = left_img.size
# normalize
processed = get_transform()
left_img = processed(left_img).numpy()
right_img = processed(right_img).numpy()
# pad to size 1248x384
top_pad = 384 - h
right_pad = 1248 - w
assert top_pad > 0 and right_pad > 0
# pad images
left_img = np.lib.pad(left_img, ((0, 0), (top_pad, 0), (0, right_pad)), mode='edge')
right_img = np.lib.pad(right_img, ((0, 0), (top_pad, 0), (0, right_pad)), mode='edge')
# pad disparity gt
if self.disp_occ_filenames is not None:
# assert len(self.disp_occ_filenames.shape) == 2
disp_occ = np.lib.pad(disp_occ, ((top_pad, 0), (0, right_pad)), mode='constant', constant_values=0)
if self.disp_occ_filenames is not None:
return {"left": left_img,
"right": right_img,
"left_disp": disp_occ,
"top_pad": top_pad,
"right_pad": right_pad
}
else:
return {"left": left_img,
"right": right_img,
"top_pad": top_pad,
"right_pad": right_pad,
"left_filename": self.left_filenames[index],
"right_filename": self.right_filenames[index]
}
def __getitem__(self, index):
left_img = self.load_image(
os.path.join(self.datapath, self.left_filenames[index]))
right_img = self.load_image(
os.path.join(self.datapath, self.right_filenames[index]))
if self.mask_filenames:
mask = self.load_image(
os.path.join(self.datapath, self.mask_filenames[index]))
else:
mask = None
if self.disp_filenames: # has disparity ground truth
disparity = self.load_disp(
os.path.join(self.datapath, self.disp_filenames[index]))
if self.mask_filenames:
mask = np.asarray(mask)
temp = mask > 0
disparity = disparity * temp
else:
disparity = None
if self.training:
#rgb2gray
left_img = left_img.convert('L')
right_img = right_img.convert('L')
w, h = left_img.size
crop_w, crop_h = 512, 256
x1 = random.randint(0, w - crop_w)
y1 = random.randint(0, h - crop_h)
# random crop
left_img = left_img.crop((x1, y1, x1 + crop_w, y1 + crop_h))
right_img = right_img.crop((x1, y1, x1 + crop_w, y1 + crop_h))
disparity = disparity[y1:y1 + crop_h, x1:x1 + crop_w]
left_img = np.ascontiguousarray(left_img, dtype=np.float32)
right_img = np.ascontiguousarray(right_img, dtype=np.float32)
# to tensor, normalize
preprocess = get_transform()
left_img = preprocess(left_img)
right_img = preprocess(right_img)
disparity = np.expand_dims(disparity, 0)
# return [left_img,right_img],-disparity
return {
"left": left_img,
"right": right_img,
"disparity": disparity
}
else:
w, h = left_img.size
top_pad = 384 - h
right_pad = 1280 - w
assert top_pad > 0 and right_pad > 0
left_img = np.ascontiguousarray(left_img, dtype=np.float32)
right_img = np.ascontiguousarray(right_img, dtype=np.float32)
left_img = np.lib.pad(
left_img,
((top_pad, 0), (0, right_pad)),
mode='symmetric',
)
right_img = np.lib.pad(right_img, ((top_pad, 0), (0, right_pad)),
mode='symmetric')
disparity = np.lib.pad(disparity, ((top_pad, 0), (0, right_pad)),
mode='constant',
constant_values=0)
preprocess = get_transform()
left_img = preprocess(left_img)
right_img = preprocess(right_img)
disparity = np.expand_dims(disparity, 0)
# return [left_img,right_img],-disparity
return {
"left": left_img,
"right": right_img,
"disparity": disparity
}
