def save_depth():
# dataset, dataloader
MVSDataset = find_dataset_def(args.dataset)
# test_dataset = MVSDataset(args.testpath, args.testlist, "test", 5, args.numdepth, args.interval_scale)
test_dataset = MVSDataset(args.testpath)
TestImgLoader = DataLoader(test_dataset,
args.batch_size,
shuffle=False,
num_workers=4,
drop_last=False)
# model
model = EdgeFlow()
model = nn.DataParallel(model)
model.cuda()
# load checkpoint file specified by args.loadckpt
print("loading model {}".format(args.loadckpt))
state_dict = torch.load(args.loadckpt)
model.load_state_dict(state_dict['model'])
model.eval()
with torch.no_grad():
for batch_idx, sample in enumerate(TestImgLoader):
sample_cuda = tocuda(sample)
# print(sample_cuda)
depth, confidence, _ = model(sample_cuda["imgs"],
sample_cuda["proj_matrices"],
sample_cuda["depth_values"])
depth = tensor2numpy(depth)
confidence = tensor2numpy(confidence)
del sample_cuda
torch.cuda.empty_cache()
print('Iter {}/{}'.format(batch_idx, len(TestImgLoader)))
filenames = sample["filename"]
# save depth maps and confidence maps
for filename, depth_est, photometric_confidence in zip(
filenames, depth["{}x".format(args.stage)],
confidence["{}x".format(args.stage)]):
depth_filename = os.path.join(
args.outdir,
filename.format('depth_est_{}x'.format(args.stage),
'.pfm'))
confidence_filename = os.path.join(
args.outdir,
filename.format('confidence_{}x'.format(args.stage),
'.pfm'))
# print(photometric_confidence.shape)
# mask = sample_cuda["mask"]
os.makedirs(depth_filename.rsplit('/', 1)[0], exist_ok=True)
os.makedirs(confidence_filename.rsplit('/', 1)[0],
exist_ok=True)
# save depth maps
save_pfm(depth_filename, depth_est)
# save confidence maps
save_pfm(confidence_filename, photometric_confidence)
def save_depth():
# dataset, dataloader
MVSDataset = find_dataset_def(args.dataset)
test_dataset = MVSDataset(args.testpath, args.n_views)
TestImgLoader = DataLoader(test_dataset,
args.batch_size,
shuffle=False,
num_workers=4,
drop_last=False)
# model
model = PatchmatchNet(
patchmatch_interval_scale=args.patchmatch_interval_scale,
propagation_range=args.patchmatch_range,
patchmatch_iteration=args.patchmatch_iteration,
patchmatch_num_sample=args.patchmatch_num_sample,
propagate_neighbors=args.propagate_neighbors,
evaluate_neighbors=args.evaluate_neighbors)
model = nn.DataParallel(model)
model.cuda()
# load checkpoint file specified by args.loadckpt
print("loading model {}".format(args.loadckpt))
state_dict = torch.load(args.loadckpt)
model.load_state_dict(state_dict['model'])
model.eval()
with torch.no_grad():
for batch_idx, sample in enumerate(TestImgLoader):
start_time = time.time()
sample_cuda = tocuda(sample)
outputs = model(sample_cuda["imgs"], sample_cuda["proj_matrices"],
sample_cuda["depth_min"], sample_cuda["depth_max"])
outputs = tensor2numpy(outputs)
del sample_cuda
print('Iter {}/{}, time = {:.3f}'.format(batch_idx,
len(TestImgLoader),
time.time() - start_time))
filenames = sample["filename"]
# save depth maps and confidence maps
for filename, depth_est, photometric_confidence in zip(
filenames, outputs["refined_depth"]['stage_0'],
outputs["photometric_confidence"]):
depth_filename = os.path.join(
args.outdir, filename.format('depth_est', '.pfm'))
confidence_filename = os.path.join(
args.outdir, filename.format('confidence', '.pfm'))
os.makedirs(depth_filename.rsplit('/', 1)[0], exist_ok=True)
os.makedirs(confidence_filename.rsplit('/', 1)[0],
exist_ok=True)
# save depth maps
depth_est = np.squeeze(depth_est, 0)
save_pfm(depth_filename, depth_est)
# save confidence maps
save_pfm(confidence_filename, photometric_confidence)
def save_depth():
# dataset, dataloader
MVSDataset = find_dataset_def(args.dataset)
test_dataset = MVSDataset(args.testpath, args.testlist, "test", 3,
args.numdepth, args.interval_scale)
TestImgLoader = DataLoader(test_dataset,
args.batch_size,
shuffle=False,
num_workers=4,
drop_last=False)
# model
model = MVSNet(refine=False)
model = nn.DataParallel(model)
model.cuda()
# load checkpoint file specified by args.loadckpt
print("loading model {}".format(args.loadckpt))
state_dict = torch.load(args.loadckpt)
model.load_state_dict(state_dict['model'])
model.eval()
with torch.no_grad():
for batch_idx, sample in enumerate(TestImgLoader):
sample_cuda = tocuda(sample)
print("begin forward")
outputs = model(sample_cuda["imgs"], sample_cuda["proj_matrices"],
sample_cuda["depth_values"])
print("forward over")
outputs = tensor2numpy(outputs)
del sample_cuda
print('Iter {}/{}'.format(batch_idx, len(TestImgLoader)))
filenames = sample["filename"]
# save depth maps and confidence maps
for filename, depth_est, photometric_confidence in zip(
filenames, outputs["depth"],
outputs["photometric_confidence"]):
depth_filename = os.path.join(
args.outdir, filename.format('depth_est', '.pfm'))
confidence_filename = os.path.join(
args.outdir, filename.format('confidence', '.pfm'))
os.makedirs(depth_filename.rsplit('/', 1)[0], exist_ok=True)
os.makedirs(confidence_filename.rsplit('/', 1)[0],
exist_ok=True)
# save depth maps
print("save .pfm")
save_pfm(depth_filename, depth_est)
# save confidence maps
save_pfm(confidence_filename, photometric_confidence)
def probability_filter(dense_folder, prob_threshold):
image_folder = os.path.join(dense_folder, 'images')
# convert cameras
image_names = os.listdir(image_folder)
for image_name in image_names:
image_prefix = os.path.splitext(image_name)[0]
init_depth_map_path = os.path.join(dense_folder, "depth_est",
image_prefix + '.pfm')
prob_map_path = os.path.join(dense_folder, "confidence",
image_prefix + '.pfm')
out_depth_map_path = os.path.join(dense_folder, "depth_est",
image_prefix + '_prob_filtered.pfm')
depth_map, _ = read_pfm(init_depth_map_path)
prob_map, _ = read_pfm(prob_map_path)
depth_map[prob_map < prob_threshold] = 0
save_pfm(out_depth_map_path, depth_map)
def save_scene_depth(testlist):
# dataset, dataloader
MVSDataset = find_dataset_def(args.dataset)
test_dataset = MVSDataset(args.testpath,
testlist,
"test",
args.num_view,
args.numdepth,
Interval_Scale,
max_h=args.max_h,
max_w=args.max_w,
fix_res=args.fix_res)
TestImgLoader = DataLoader(test_dataset,
args.batch_size,
shuffle=False,
num_workers=4,
drop_last=False)
# model
model = DDR_Net(
refine=False,
ndepths=[int(nd) for nd in args.ndepths.split(",") if nd],
depth_interals_ratio=[
float(d_i) for d_i in args.depth_inter_r.split(",") if d_i
],
share_cr=args.share_cr,
cr_base_chs=[int(ch) for ch in args.cr_base_chs.split(",") if ch],
grad_method=args.grad_method)
# load checkpoint file specified by args.loadckpt
print("loading model {}".format(args.loadckpt))
state_dict = torch.load(args.loadckpt, map_location=torch.device("cpu"))
model.load_state_dict(state_dict['model'], strict=True)
model = nn.DataParallel(model)
model.cuda()
model.eval()
with torch.no_grad():
for batch_idx, sample in enumerate(TestImgLoader):
sample_cuda = tocuda(sample)
start_time = time.time()
outputs, refine_depth_map = model(sample_cuda["imgs"],
sample_cuda["proj_matrices"],
sample_cuda["depth_values"])
end_time = time.time()
#outputs = tensor2numpy(outputs)
depth = outputs["depth"]
#feature=outputs
photometric_confidence = outputs["photometric_confidence"]
depth = tensor2numpy(depth)
photometric_confidence = tensor2numpy(photometric_confidence)
del sample_cuda
filenames = sample["filename"]
cams = sample["proj_matrices"]["stage{}".format(num_stage)].numpy()
imgs = sample["imgs"].numpy()
print('Iter {}/{}, Time:{} Res:{}'.format(batch_idx,
len(TestImgLoader),
end_time - start_time,
imgs[0].shape))
# save depth maps and confidence maps
for filename, cam, img, depth_est, photometric_confidence in zip(filenames, cams, imgs, \
depth,photometric_confidence ):
img = img[0] #ref view
cam = cam[0] #ref cam
depth_filename = os.path.join(
args.outdir, filename.format('depth_est', '.pfm'))
confidence_filename = os.path.join(
args.outdir, filename.format('confidence', '.pfm'))
cam_filename = os.path.join(
args.outdir, filename.format('cams', '_cam.txt'))
img_filename = os.path.join(args.outdir,
filename.format('images', '.jpg'))
ply_filename = os.path.join(
args.outdir, filename.format('ply_local', '.ply'))
feature_filename = os.path.join(
args.outdir, filename.format('feature', '.jpg'))
os.makedirs(depth_filename.rsplit('/', 1)[0], exist_ok=True)
os.makedirs(confidence_filename.rsplit('/', 1)[0],
exist_ok=True)
os.makedirs(cam_filename.rsplit('/', 1)[0], exist_ok=True)
os.makedirs(img_filename.rsplit('/', 1)[0], exist_ok=True)
os.makedirs(ply_filename.rsplit('/', 1)[0], exist_ok=True)
#save depth maps
save_pfm(depth_filename, depth_est)
#save confidence maps
save_pfm(confidence_filename, photometric_confidence)
#save cams, img
write_cam(cam_filename, cam)
img = np.clip(np.transpose(img, (1, 2, 0)) * 255, 0,
255).astype(np.uint8)
img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
cv2.imwrite(img_filename, img_bgr)
# vis
# print(photometric_confidence.mean(), photometric_confidence.min(), photometric_confidence.max())
# import matplotlib.pyplot as plt
# plt.subplot(1, 3, 1)
# plt.imshow(img)
# plt.subplot(1, 3, 2)
# plt.imshow((depth_est - depth_est.min())/(depth_est.max() - depth_est.min()))
# plt.subplot(1, 3, 3)
# plt.imshow(photometric_confidence)
# plt.show()
if num_stage == 1:
downsample_img = cv2.resize(
img,
(int(img.shape[1] * 0.25), int(img.shape[0] * 0.25)))
elif num_stage == 2:
downsample_img = cv2.resize(
img,
(int(img.shape[1] * 0.5), int(img.shape[0] * 0.5)))
elif num_stage == 3:
downsample_img = img
if batch_idx % args.save_freq == 0:
generate_pointcloud(downsample_img, depth_est,
ply_filename, cam[1, :3, :3])
torch.cuda.empty_cache()
gc.collect()
def save_depth():
# dataset, dataloader
MVSDataset = find_dataset_def(args.dataset)
test_dataset = MVSDataset(args.testpath, args.testlist, "test", 5,
args.numdepth, args.interval_scale,
args.inverse_depth, args.pyramid)
TestImgLoader = DataLoader(test_dataset,
args.batch_size,
shuffle=False,
num_workers=4,
drop_last=False)
# model
if args.model == 'mvsnet':
print('use MVSNet')
model = MVSNet(refine=args.refine,
fea_net=args.fea_net,
cost_net=args.cost_net,
refine_net=args.refine_net,
origin_size=args.origin_size,
cost_aggregation=args.cost_aggregation,
dp_ratio=args.dp_ratio)
else:
print('input pre-defined model')
model = nn.DataParallel(model)
model.cuda()
# load checkpoint file specified by args.loadckpt
print("loading model {}".format(args.loadckpt))
state_dict = torch.load(args.loadckpt)
model.load_state_dict(state_dict['model'])
model.eval()
count = -1
total_time = 0
with torch.no_grad():
for batch_idx, sample in enumerate(TestImgLoader):
count += 1
print('process', sample['filename'])
sample_cuda = tocuda(sample)
print(sample_cuda["imgs"].shape)
time_s = time.time()
outputs = model(sample_cuda["imgs"], sample_cuda["proj_matrices"],
sample_cuda["depth_values"])
one_time = time.time() - time_s
total_time += one_time
print('one forward: ', one_time)
if count % 50 == 0:
print('avg time:', total_time / 50)
total_time = 0
if 'High' in args.fea_net and 'Coarse2Fine' in args.cost_net:
tmp_outputs = {}
for key, value in outputs.items():
tmp_outputs[key] = value[0]
outputs = tmp_outputs
outputs = tensor2numpy(outputs)
del sample_cuda
print('Iter {}/{}'.format(batch_idx, len(TestImgLoader)))
filenames = sample["filename"]
# save depth maps and confidence maps
for filename, depth_est, photometric_confidence in zip(
filenames, outputs["depth"],
outputs["photometric_confidence"]):
depth_filename = os.path.join(
save_dir,
filename.format('depth_est_{}'.format(args.pyramid),
'.pfm'))
confidence_filename = os.path.join(
save_dir,
filename.format('confidence_{}'.format(args.pyramid),
'.pfm'))
os.makedirs(depth_filename.rsplit('/', 1)[0], exist_ok=True)
os.makedirs(confidence_filename.rsplit('/', 1)[0],
exist_ok=True)
# save depth maps
print(depth_est.shape)
save_pfm(depth_filename, depth_est.squeeze())
# save confidence maps
save_pfm(confidence_filename, photometric_confidence.squeeze())
def multi_scale(depth_folder_down4, depth_folder_down8, depth_folder_down16,
final_folder, test_list):
depth_folder = depth_folder_down4
conf_folder = depth_folder
output_folder = depth_folder
depth2_folder = depth_folder_down8
conf2_folder = depth2_folder
mask2_folder = depth2_folder
depth3_folder = depth_folder_down16
conf3_folder = depth3_folder
mask3_folder = depth3_folder
if not os.path.exists(final_folder):
os.mkdir(final_folder)
file = open(test_list)
scans = file.readlines()
scans = [line.rstrip() for line in scans]
sum = 0
ct = 0
shape_all = 0
for scan in scans:
ct = ct + 1
depth_t = os.path.join(depth_folder, scan)
depth2_t = os.path.join(depth2_folder, scan)
depth3_t = os.path.join(depth3_folder, scan)
conf_t = os.path.join(conf_folder, scan)
conf2_t = os.path.join(conf2_folder, scan)
conf3_t = os.path.join(conf3_folder, scan)
mask2_folder_t = os.path.join(mask2_folder, scan)
mask2_folder_t = os.path.join(mask2_folder_t, 'mask')
mask3_folder_t = os.path.join(mask3_folder, scan)
mask3_folder_t = os.path.join(mask3_folder_t, 'mask')
output_t = os.path.join(output_folder, 'output')
final_t = os.path.join(final_folder, scan)
if not os.path.exists(final_t):
os.mkdir(final_t)
if not os.path.exists(output_t):
os.mkdir(output_t)
output_t = os.path.join(output_t, scan)
if not os.path.exists(output_t):
os.mkdir(output_t)
sum_t = 0
shape = 0
for i in range(49):
print('process depth ' + str(i))
depth_s = os.path.join(depth_t, 'depth_est')
final_d = os.path.join(final_t, 'depth_est')
if not os.path.exists(final_d):
os.mkdir(final_d)
final_c = os.path.join(final_t, 'confidence')
if not os.path.exists(final_c):
os.mkdir(final_c)
depth_s = os.path.join(depth_s, '%08d.pfm' % i)
depth2_s = os.path.join(depth2_t, 'depth_est')
depth2_s = os.path.join(depth2_s, '%08d.pfm' % i)
mask2_s = os.path.join(mask2_folder_t, '%08d_final.png' % i)
depth3_s = os.path.join(depth3_t, 'depth_est')
depth3_s = os.path.join(depth3_s, '%08d.pfm' % i)
mask3_s = os.path.join(mask3_folder_t, '%08d_final.png' % i)
conf_s = os.path.join(conf_t, 'confidence')
conf_s = os.path.join(conf_s, '%08d.pfm' % i)
final_d_s = os.path.join(final_d, '%08d.pfm' % i)
final_c_s = os.path.join(final_c, '%08d.pfm' % i)
conf2_s = os.path.join(conf2_t, 'confidence')
conf2_s = os.path.join(conf2_s, '%08d.pfm' % i)
conf3_s = os.path.join(conf3_t, 'confidence')
conf3_s = os.path.join(conf3_s, '%08d.pfm' % i)
output_s = os.path.join(output_t, 'error_map_%04d.png' % i)
mask2_img_t = cv2.imread(mask2_s)
mask2_img_t = cv2.cvtColor(mask2_img_t, cv2.COLOR_BGR2GRAY)
mask2_arr = np.array(mask2_img_t)
mask2_t = mask2_arr > 0
mask3_img_t = cv2.imread(mask3_s)
mask3_img_t = cv2.cvtColor(mask3_img_t, cv2.COLOR_BGR2GRAY)
mask3_arr = np.array(mask3_img_t)
mask3_arr = cv2.pyrUp(mask3_arr)
depth = read_pfm(depth_s)[0]
depth2 = read_pfm(depth2_s)[0]
depth3 = read_pfm(depth3_s)[0]
depth3 = cv2.pyrUp(depth3)
confidence = read_pfm(conf_s)[0]
confidence2 = read_pfm(conf2_s)[0]
confidence3 = read_pfm(conf3_s)[0]
confidence3 = cv2.pyrUp(confidence3)
depth2 = cv2.pyrUp(depth2)
zero_depth = np.zeros([8, 400])
depth2 = np.r_[depth2, zero_depth]
confidence2 = cv2.pyrUp(confidence2)
zero_conf = np.zeros([8, 400])
confidence2 = np.r_[confidence2, zero_conf]
mask2_arr = cv2.pyrUp(mask2_arr)
zero_depth = np.zeros([8, 400])
mask2_arr = np.r_[mask2_arr, zero_depth]
mask3 = (0.9 < confidence2) & (confidence < 0.5) & (mask2_arr > 0)
depth[mask3] = depth2[mask3]
confidence[mask3] = confidence2[mask3]
save_pfm(final_d_s, depth)
save_pfm(final_c_s, confidence)
