def __init__(self, train_file, mode='train', n_samples=None):
self.mode = mode
self.train_file = open(train_file,'r')
self.image_lists = []
for line in self.train_file:
line = line.strip('\n')
self.image_lists.append(line)
self.data_root = 'data/kitti'
#left image
self.left_image = cv2.imread('data/driving/left/0350.png',0)
self.left_image = self.left_image.astype(np.float32)
self.max_pixel = np.max(self.left_image)
#use disp image .pfm
self.disp_image, _ = readPFM('data/driving/disp/0350.pfm')
self.disp_image = self.disp_image.astype(np.int32)
# #use disp image .png
# disp_path = self.data_root + '/disp_occ_0/' + image_name + '.png'
# disp_image = cv2.imread(disp_path,0)
# use the ct_cost from calculate_ct_cost in local file
self.d_cost = np.load('data/driving/disp/0350.npy')
self.d_cost = np.array(self.d_cost)
self.d_cost = torch.tensor(self.d_cost)
self.d_cost = self.d_cost.type(torch.FloatTensor)
#use the cost volume geted from mc-cnn https://github.com/jzbontar/mc-cnn
# cost_path = self.data_root + '/leftbin/' + image_name + '.bin'
# d_cost = np.memmap(cost_path, dtype=np.float32, shape=(1, 70, img_h, img_w))
# p = np.isnan(d_cost)
# d_cost[p] = 0.0
# d_cost = torch.tensor(d_cost)
# d_cost = d_cost.squeeze(0)
# d_cost = d_cost.type(torch.FloatTensor) # [70,h,w]
self.img_h , self.img_w = self.left_image.shape
self.valid_points = []
#get the valid points list
for i in range(300, self.img_w-200):
for j in range(150, self.img_h -150):
disp_val = self.disp_image[j,i]
if(disp_val>5 and disp_val < 69):
self.valid_points.append([i,j,disp_val])
random.shuffle(self.valid_points)
self.valid_points = random.sample(self.valid_points, int(len(self.valid_points)/20))
print(self.d_cost.size())
print(self.img_h,self.img_w)
def disparity_loader(path):
path_prefix = path.split('.')[0]
# print(path_prefix)
path1 = path_prefix + '_exception_assign_minus_1.npy'
path2 = path_prefix + '.npy'
path3 = path_prefix + '.pfm'
import os.path as ospath
if ospath.exists(path1):
return np.load(path1)
else:
# from readpfm import readPFMreadPFM
from readpfm import readPFM
data, _ = readPFM(path3)
np.save(path2, data)
#for i in range(data.shape[0]):
# for j in range(data.shape[1]):
# if j - data[i][j] < 0:
# data[i][j] = -1
np.save(path1, data)
return data
def disparity_loader(path):
return rp.readPFM(path)
def disp_loader(self, path):
disp, scale = readPFM(path)
disp = np.ascontiguousarray(disp, dtype=np.float32)
disp[np.inf == disp] = np.nan
return disp
import argparse
parser = argparse.ArgumentParser(description='pfm visualization')
parser.add_argument('--input_dir',
type=str,
default='./dataset/lr_disp',
help='depth_path')
parser.add_argument('--output_dir',
type=str,
default='./result',
help='outputdir')
args = parser.parse_args()
if os.path.isdir(args.input_dir):
test_pfm = os.listdir(args.input_dir)
test_pfm.sort()
pfm_path = [os.path.join(args.input_dir, d) for d in test_pfm]
else:
pfm_path = [args.input_dir]
print("there are %d files" % (len(pfm_path)))
basenames = [os.path.basename(f) for f in pfm_path]
basenames = [os.path.splitext(f)[0] for f in basenames]
for num, path in enumerate(pfm_path):
output = readPFM(path)
plt.imsave(str(os.path.join(args.output_dir, basenames[num]) + "_vis.png"),
output,
cmap='plasma')
print("finished %d" % (num + 1))