def KT15FalseColorDisp(batch_inp, max_disp = 256):
"""
args:
batch_inp: in shape [N,C,H,W]
"""
"""Convert a Tensor to numpy image."""
#print ("[??]batch_inp shape = ", batch_inp.shape)
#batch_inp = batch_inp.cpu().numpy().transpose((0,2,3,1))# change to [N,H,W,C]
batch_inp = batch_inp.cpu().numpy() # in shape [N,C=1,H,W]
tmp = []
for i in range(batch_inp.shape[0]):
tmp.append(KT15FalseClr.writeKT15FalseColor(
np.squeeze(batch_inp[i],0),
max_disp)[:,:,::-1] # RGB to BGR for tensorboard image write
)
return np.stack(tmp, axis = 0)
def test(self):
self.model.eval()
file_path = self.args.data_path
file_list_txt = self.args.test_list
f = open(file_list_txt, 'r')
if self.virtual_kitti2:
filelist = get_virtual_kitti2_filelist(file_list_txt)
else:
filelist = [l.rstrip() for l in f.readlines() if not l.rstrip().startswith('#')]
crop_width = self.args.crop_width
crop_height = self.args.crop_height
if not os.path.exists(self.args.resultDir):
os.makedirs(self.args.resultDir)
print ('makedirs {}'.format(self.args.resultDir))
avg_err = 0
avg_rate1 = 0
avg_rate3 = 0
for index in range(len(filelist)):
#for index in range(1):
current_file = filelist[index]
if self.kitti2015:
leftname = pjoin(file_path, 'image_0/' + current_file)
if index % 20 == 0:
print ("limg: {}".format(leftname))
rightname = pjoin(file_path, 'image_1/' + current_file)
dispname = pjoin(file_path, 'disp_occ_0_pfm/' + current_file[0:-4] + '.pfm')
if os.path.isfile(dispname):
dispGT = pfm.readPFM(dispname)
dispGT[dispGT == np.inf] = .0
else:
dispGT= None
savename = pjoin(self.args.resultDir, current_file[0:-4] + '.pfm')
elif self.kitti2012:
#leftname = pjoin(file_path, 'image_0/' + current_file)
leftname = pjoin(file_path, 'colored_0/' + current_file)
rightname = pjoin(file_path, 'colored_1/' + current_file)
dispname = pjoin(file_path, 'disp_occ_0_pfm/' + current_file[0:-4] + '.pfm')
if os.path.isfile(dispname):
dispGT = pfm.readPFM(dispname)
dispGT[dispGT == np.inf] = .0
else:
dispGT= None
savename = pjoin(self.args.resultDir, current_file[0:-4] + '.pfm')
#disp = Image.open(dispname)
#disp = np.asarray(disp) / 256.0
elif self.virtual_kitti2:
A = current_file
# e.g., /media/ccjData2/datasets/Virtual-KITTI-V2/vkitti_2.0.3_rgb/Scene01/15-deg-left/frames/rgb/Camera_0/rgb_00001.jpg
leftname = pjoin(file_path, "vkitti_2.0.3_rgb/" + A)
rightname = pjoin(file_path, "vkitti_2.0.3_rgb/" + A[:-22] + 'Camera_1/' + A[-13:])
#load depth GT and change it to disparity GT:
depth_png_filename = pjoin(file_path, "vkitti_2.0.3_depth/" + A[:-26] + 'depth/Camera_0/depth_' + A[-9:-4] + ".png")
#print ("imgl = ", leftname, ", imgr = ", rightname, ", depth_left = ", depth_png_filename)
#NOTE: The depth map in centimeters can be directly loaded
depth_left = cv2.imread(depth_png_filename, cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH)
height, width = depth_left.shape[:2]
# Intrinsi: f_x = f_y = 725.0087
# offset(i.e., distance between stereo views): B = 0.532725 m = 53.2725 cm
B = 53.2725 # in centimeters;
f = 725.0087 # in pixels
# set zero as a invalid disparity value;
dispGT = np.zeros([height, width], 'float32')
mask = depth_left > 0
dispGT[mask] = f*B/ depth_left[mask] # d = fB/z
#pfm.show(dispGT, title='dispGT')
savename = pjoin(self.args.resultDir, '%04d.pfm'%(index))
elif self.eth3d or self.middlebury:
leftname = pjoin(file_path, current_file + '/im0.png')
rightname = pjoin(file_path, current_file + '/im1.png')
print ("limg: {}".format(leftname))
H, W = cv2.imread(leftname).shape[:2]
dispname = pjoin(file_path, current_file + '/disp0GT.pfm')
if os.path.isfile(dispname):
dispGT = pfm.readPFM(dispname)
dispGT[dispGT == np.inf] = .0
else:
dispGT= None
savename = pjoin(self.args.resultDir, current_file + '.pfm')
# due to eth3d images have diverse image size;
# so we process each of them seperately;
ds = self.ds # e.g., == 64, encoder sacle;
crop_width = W + (ds - W % ds) % ds
crop_height = H + (ds - H % ds) % ds
print ("input HxW=%dx%d, padding size=%dx%d" %(H,W,crop_height, crop_width))
else:
A = current_file
leftname = pjoin(file_path, A)
rightname = pjoin(file_path, A[:-13] + 'right/' + A[len(A)-8:])
# check disparity GT exists or not!!!
pos = A.find('/')
tmp_len = len('frames_finalpass')
dispname = pjoin(file_path, A[0:pos] + '/disparity' + A[pos+1+tmp_len:-4] + '.pfm')
#print ("[****] ldisp: {}".format(dispname))
if os.path.isfile(dispname):
dispGT = pfm.readPFM(dispname)
dispGT[dispGT == np.inf] = .0
else:
dispGT= None
savename = pjoin(self.args.resultDir, str(index) + '.pfm')
input1, input2, height, width = test_transform(
load_test_data(leftname, rightname, is_data_augment=self.is_data_augment),
crop_height, crop_width)
if self.cuda:
input1 = input1.cuda()
input2 = input2.cuda()
with torch.no_grad():
prediction = self.model(input1, input2)
disp = prediction.cpu().detach().numpy() # in full size
if height <= crop_height and width <= crop_width:
disp = disp[0, crop_height - height: crop_height, crop_width-width: crop_width]
else:
disp = disp[0, :, :]
# for KT benchmark submission;
if dispGT is None:
tmp_savename = savename[:-4] + '.png'
print ("savd ", tmp_savename)
skimage.io.imsave(tmp_savename, (disp * 256).astype('uint16'))
#sys.exit()
if any([self.kitti2015, self.kitti2012, self.eth3d, self.middlebury, index %250 == 0]):
pfm.save(savename, disp)
#print ('saved ', savename)
if 0 and dispGT is not None:
left = np.asarray(Image.open(leftname))[:,:,:3].astype(np.float32)
right = np.asarray(Image.open(rightname))[:,:,:3].astype(np.float32)
#print ("[???]", left.shape)
pfm.save(savename[:-4] + '-iml.pfm', left)
pfm.save(savename[:-4] + '-imr.pfm', right)
pfm.save(savename[:-4] + '-gt.pfm', dispGT.astype(np.float32))
if dispGT is not None:
#error, rate1, rate3 = get_epe_rate2(dispGT, disp, self.args.max_disp, threshold=1.0, threshold2=3.0)
error, rate1, rate3 = get_epe_rate(dispGT, disp, threshold=1.0, threshold2=3.0)
avg_err += error
avg_rate1 += rate1
avg_rate3 += rate3
if index % 250 == 0:
print("===> Frame {}: ".format(index) + leftname + " ==> EPE Error: {:.4f}, Bad-{:.1f} Error: {:.4f}, Bad-{:.1f} Error: {:.4f}".format(
error, 1.0, rate1, 3.0, rate3))
# save kt15 color
#if self.kitti2015:
if any([self.kitti2015, self.kitti2012, self.eth3d, self.middlebury, index%250==0]):
""" disp """
tmp_dir = pjoin(self.args.resultDir, "dispColor")
if not os.path.exists(tmp_dir):
os.makedirs(tmp_dir)
if self.eth3d or self.middlebury:
tmp_dispname = pjoin(tmp_dir, current_file + '.png')
elif self.kitti2012 or self.kitti2015:
tmp_dispname = pjoin(tmp_dir, current_file[0:-4] + '.png')
else:
tmp_dispname = pjoin(tmp_dir, '%04d.png'%(index))
cv2.imwrite(tmp_dispname,
KT15FalseClr.writeKT15FalseColor(np.ascontiguousarray(disp)).astype(np.uint8)[:,:,::-1])
if index < 1:
print ('saved ', tmp_dispname)
if dispGT is not None: #If KT benchmark submission, then No dispGT;
""" err-disp """
tmp_dir = pjoin(self.args.resultDir, "errDispColor")
if not os.path.exists(tmp_dir):
os.makedirs(tmp_dir)
if self.eth3d or self.middlebury:
tmp_errdispname = pjoin(tmp_dir, current_file + '.png')
elif self.kitti2015 or self.kitti2012:
tmp_errdispname = pjoin(tmp_dir, current_file[0:-4] + '.png')
else:
tmp_errdispname = pjoin(tmp_dir, '%04d.png'%(index))
cv2.imwrite(tmp_errdispname,
KT15LogClr.writeKT15ErrorDispLogColor(np.ascontiguousarray(disp), np.ascontiguousarray(dispGT)).astype(np.uint8)[:,:,::-1])
if index < 1:
print ('saved ', tmp_errdispname)
if dispGT is not None:
avg_err /= len(filelist)
avg_rate1 /= len(filelist)
avg_rate3 /= len(filelist)
print("===> Total {} Frames ==> AVG EPE Error: {:.4f}, AVG Bad-{:.1f} Error: {:.4f}, AVG Bad-{:.1f} Error: {:.4f}".format(
len(filelist), avg_err, 1.0, avg_rate1, 3.0, avg_rate3))
""" save as csv file, Excel file format """
csv_file = os.path.join(self.args.resultDir, 'bad-err.csv')
print ("write ", csv_file, "\n")
timeStamp = datetime.now().strftime('%Y-%m-%d_%H:%M:%S')
messg = timeStamp + ',{},bad-1.0,{:.4f},bad-3.0,{:.4f},epe,{:.4f},fileDir={},for log,{:.3f}(epe); {:.3f}%(bad1); {:.3f}%(bad3)\n'.format(
self.data_type_str, avg_rate1, avg_rate3, avg_err,
self.args.resultDir,
avg_err, avg_rate1*100.0, avg_rate3*100.0)
with open( csv_file, 'w') as fwrite:
fwrite.write(messg)
if any([opt.kitti2015, opt.kitti2012, index % 250 == 0]):
pfm.save(savename, disp)
""" disp """
tmp_dir = pjoin(opt.resultDir, "dispColor")
if not os.path.exists(tmp_dir):
os.makedirs(tmp_dir)
# save kt15 color
if opt.kitti2015 or opt.kitti2012:
tmp_dispname = pjoin(tmp_dir, current_file[0:-4] + '.png')
else:
tmp_dispname = pjoin(tmp_dir, '%04d.png' % (index))
cv2.imwrite(
tmp_dispname,
KT15FalseClr.writeKT15FalseColor(
np.ascontiguousarray(disp)).astype(np.uint8)[:, :, ::-1])
print("save ", tmp_dispname)
if dispGT is not None: #If KT benchmark submission, then No dispGT;
""" err-disp """
tmp_dir = pjoin(opt.resultDir, "errDispColor")
if not os.path.exists(tmp_dir):
os.makedirs(tmp_dir)
if opt.kitti2015 or opt.kitti2012:
tmp_errdispname = pjoin(tmp_dir,
current_file[0:-4] + '.png')
else:
tmp_errdispname = pjoin(tmp_dir, '%04d.png' % (index))
cv2.imwrite(
tmp_errdispname,
def test(self, py=None):
self.model.eval()
if not os.path.exists(self.args.resultDir):
os.makedirs(self.args.resultDir)
print('makedirs {}'.format(self.args.resultDir))
dispScale = 1.0
avg_err = .0
avg_rate = .0
img_num = self.test_loader_len
print("[****] To test %d images !!!" % img_num)
for iteration, batch_data in enumerate(self.testing_data_loader):
batch_size = batch_data[0].size()[0]
assert batch_size == 1
features = batch_data[0]
print("[???] features size:", features.size())
height = batch_data[1][0].item()
width = batch_data[2][0].item()
crop_height = batch_data[3][0].item()
crop_width = batch_data[4][0].item()
current_file = batch_data[5][0]
disp_name = batch_data[6][0]
print(height, crop_height, current_file, disp_name)
if os.path.isfile(disp_name):
dispGT = pfm.readPFM(disp_name)
dispGT[dispGT == np.inf] = .0
else:
dispGT = None
if self.kitti2015 or self.kitti2012:
savename = pjoin(self.args.resultDir,
current_file[0:-4] + '.pfm')
elif self.eth3d or self.middlebury:
savename = pjoin(self.args.resultDir, current_file + '.pfm')
else: # scene flow dataset
savename = pjoin(self.args.resultDir, str(iteration) + '.pfm')
if self.cuda:
features = features.cuda()
with torch.no_grad():
if self.model_name == 'MS-GCNet':
disp = self.model(features)
elif self.model_name == 'MS-PSMNet':
disp = self.model(features)
else:
raise Exception(
"No suitable model found ... Wrong name: {}".format(
self.model_name))
#about memory
del features
disp = disp.cpu().detach().numpy() * dispScale
if height <= crop_height and width <= crop_width:
#disp = disp[0, crop_height - height: crop_height, crop_width-width: crop_width]
disp = disp[0, crop_height - height:crop_height, 0:width]
else:
disp = disp[0, :, :]
#save to uint16 png files
#skimage.io.imsave(savename, (disp * 256).astype('uint16'))
if any([
self.kitti2015, self.kitti2012, self.eth3d,
self.middlebury, iteration % 50 == 0
]):
pfm.save(savename, disp)
#print ('saved ', savename)
if dispGT is not None:
if self.eth3d:
self.args.threshold = 1.0
elif self.middlebury:
self.args.threshold = 1.0 # for trainingH;
elif self.kitti2012 or self.kitti2015:
self.args.threshold = 3.0
else: # Scene Flow
self.args.threshold = 1.0
error, rate = get_epe_rate(dispGT, disp, self.max_disp,
self.args.threshold)
avg_err += error
avg_rate += rate
if iteration % 200 == 0:
message_info = "===> Frame {}: ".format(
iteration
) + current_file + " ==> EPE Error: {:.4f}, Bad-{:.1f} Error: {:.4f}".format(
error, self.args.threshold, rate)
""" adapted from Mateo's code """
if py is not None:
gc.collect()
memoryUse = py.memory_info(
)[0] / 2.**20 # memory use in MB...I think
message_info += ', memeory: {:.2f} MB'.format(
memoryUse)
print(message_info)
# save kt15 color
#if self.kitti2015:
if any(
[self.kitti2015, self.kitti2012, self.eth3d, self.middlebury]):
""" disp """
tmp_dir = pjoin(self.args.resultDir, "dispColor")
if not os.path.exists(tmp_dir):
os.makedirs(tmp_dir)
tmp_dispname = pjoin(tmp_dir, current_file[0:-4] + '.png')
cv2.imwrite(
tmp_dispname,
KT15FalseClr.writeKT15FalseColor(
np.ascontiguousarray(disp)).astype(
np.uint8)[:, :, ::-1])
if iteration % 50 == 0:
print('saved ', tmp_dispname)
if dispGT is not None: #If KT benchmark submission, then No dispGT;
""" err-disp """
tmp_dir = pjoin(self.args.resultDir, "errDispColor")
if not os.path.exists(tmp_dir):
os.makedirs(tmp_dir)
tmp_errdispname = pjoin(tmp_dir,
current_file[0:-4] + '.png')
cv2.imwrite(
tmp_errdispname,
KT15LogClr.writeKT15ErrorDispLogColor(
np.ascontiguousarray(disp),
np.ascontiguousarray(dispGT)).astype(
np.uint8)[:, :, ::-1])
if iteration % 50 == 0:
print('saved ', tmp_errdispname)
if dispGT is not None:
avg_err = avg_err / img_num
avg_rate = avg_rate / img_num
print(
"===> Total {} Frames ==> AVG EPE Error: {:.4f}, AVG Bad-{:.1f} Error: {:.4f}"
.format(img_num, avg_err, self.args.threshold, avg_rate))
print('{} testing finished!'.format(self.model_name))
def test(self):
self.model.eval()
file_path = self.args.data_path
file_list = self.args.test_list
f = open(file_list, 'r')
filelist = [l.rstrip() for l in f.readlines()]
crop_width = self.args.crop_width
crop_height = self.args.crop_height
if not os.path.exists(self.args.resultDir):
os.makedirs(self.args.resultDir)
print('makedirs {}'.format(self.args.resultDir))
avg_err = 0
avg_rate = 0
for index in range(len(filelist)):
#for index in range(1):
current_file = filelist[index]
if self.kitti2015:
leftname = pjoin(file_path, 'image_0/' + current_file)
if index % 20 == 0:
print("limg: {}".format(leftname))
rightname = pjoin(file_path, 'image_1/' + current_file)
dispname = pjoin(
file_path, 'disp_occ_0_pfm/' + current_file[0:-4] + '.pfm')
if os.path.isfile(dispname):
dispGT = pfm.readPFM(dispname)
dispGT[dispGT == np.inf] = .0
else:
dispGT = None
savename = pjoin(self.args.resultDir,
current_file[0:-4] + '.pfm')
elif self.kitti2012:
#leftname = pjoin(file_path, 'image_0/' + current_file)
leftname = pjoin(file_path, 'colored_0/' + current_file)
rightname = pjoin(file_path, 'colored_1/' + current_file)
dispname = pjoin(
file_path, 'disp_occ_0_pfm/' + current_file[0:-4] + '.pfm')
if os.path.isfile(dispname):
dispGT = pfm.readPFM(dispname)
dispGT[dispGT == np.inf] = .0
else:
dispGT = None
savename = pjoin(self.args.resultDir,
current_file[0:-4] + '.pfm')
#disp = Image.open(dispname)
#disp = np.asarray(disp) / 256.0
else:
A = current_file
leftname = pjoin(file_path, A)
rightname = pjoin(file_path,
A[:-13] + 'right/' + A[len(A) - 8:])
# check disparity GT exists or not!!!
pos = A.find('/')
tmp_len = len('frames_finalpass')
dispname = pjoin(
file_path,
A[0:pos] + '/disparity' + A[pos + 1 + tmp_len:-4] + '.pfm')
#print ("[****] ldisp: {}".format(dispname))
if os.path.isfile(dispname):
dispGT = pfm.readPFM(dispname)
dispGT[dispGT == np.inf] = .0
else:
dispGT = None
savename = pjoin(self.args.resultDir, str(index) + '.pfm')
input1, input2, height, width = test_transform(
load_test_data(leftname, rightname), crop_height, crop_width)
if self.cuda:
input1 = input1.cuda()
input2 = input2.cuda()
with torch.no_grad():
prediction1, prediction = self.model(input1, input2)
disp = prediction.cpu().detach().numpy() # in full size
disp1 = prediction1.cpu().detach().numpy() # in half size
if height <= crop_height and width <= crop_width:
disp = disp[0, crop_height - height:crop_height,
crop_width - width:crop_width]
disp1 = disp1[0,
crop_height // 2 - height // 2:crop_height // 2,
crop_width // 2 - width // 2:crop_width // 2]
else:
disp = disp[0, :, :]
disp1 = disp1[0, :, :]
#skimage.io.imsave(savename, (disp * 256).astype('uint16'))
if self.kitti2015 or self.kitti2012 or index % 50 == 0:
pfm.save(savename, disp)
#print ('saved ', savename)
if 0 and dispGT is not None:
left = np.asarray(Image.open(leftname))[:, :, :3].astype(
np.float32)
right = np.asarray(Image.open(rightname))[:, :, :3].astype(
np.float32)
#print ("[???]", left.shape)
pfm.save(savename[:-4] + '-iml.pfm', left)
pfm.save(savename[:-4] + '-imr.pfm', right)
pfm.save(savename[:-4] + '-gt.pfm',
dispGT.astype(np.float32))
pfm.save(savename[:-4] + '-half.pfm', disp1)
if dispGT is not None:
error, rate = get_epe_rate(dispGT, disp, self.args.max_disp,
self.args.threshold)
avg_err += error
avg_rate += rate
if index % 20 == 0:
print("===> Frame {}: ".format(index) + leftname +
" ==> EPE Error: {:.4f}, Bad-{:.1f} Error: {:.4f}".
format(error, self.args.threshold, rate))
# save kt15 color
if self.kitti2015:
""" disp """
tmp_dir = pjoin(self.args.resultDir, "dispColor")
if not os.path.exists(tmp_dir):
os.makedirs(tmp_dir)
tmp_dispname = pjoin(tmp_dir, current_file[0:-4] + '.png')
cv2.imwrite(
tmp_dispname,
KT15FalseClr.writeKT15FalseColor(
np.ascontiguousarray(disp)).astype(
np.uint8)[:, :, ::-1])
if index % 20 == 0:
print('saved ', tmp_dispname)
if dispGT is not None: #If KT benchmark submission, then No dispGT;
""" err-disp """
tmp_dir = pjoin(self.args.resultDir, "errDispColor")
if not os.path.exists(tmp_dir):
os.makedirs(tmp_dir)
tmp_errdispname = pjoin(tmp_dir,
current_file[0:-4] + '.png')
cv2.imwrite(
tmp_errdispname,
KT15LogClr.writeKT15ErrorDispLogColor(
np.ascontiguousarray(disp),
np.ascontiguousarray(dispGT)).astype(
np.uint8)[:, :, ::-1])
if index % 20 == 0:
print('saved ', tmp_errdispname)
if dispGT is not None:
avg_err = avg_err / len(filelist)
avg_rate = avg_rate / len(filelist)
print(
"===> Total {} Frames ==> AVG EPE Error: {:.4f}, AVG Bad-{:.1f} Error: {:.4f}"
.format(len(filelist), avg_err, self.args.threshold, avg_rate))