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",
5,
args.numdepth,
interval_scale=args.interval_scale)
TestImgLoader = DataLoader(test_dataset,
args.batch_size,
shuffle=False,
num_workers=4,
drop_last=False)
# model
model = MVSNet(refine=False, upsample=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)
outputs = model(sample_cuda["imgs"], sample_cuda["proj_matrices"],
sample_cuda["depth_values"])
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', '.npy'))
confidence_filename = os.path.join(
args.outdir, filename.format('confidence', '.npy'))
os.makedirs(depth_filename.rsplit('/', 1)[0], exist_ok=True)
os.makedirs(confidence_filename.rsplit('/', 1)[0],
exist_ok=True)
# save depth maps
np.save(depth_filename, depth_est)
# save_pfm(depth_filename, depth_est)
# save confidence maps
np.save(confidence_filename, photometric_confidence)
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()
# p: probability volume [B, D, H, W]
# depth_values: discrete depth values [B, D]
def depth_regression(p, depth_values):
depth_values = depth_values.view(*depth_values.shape, 1, 1)
depth = torch.sum(p * depth_values, 1)
return depth
if __name__ == "__main__":
# some testing code, just IGNORE it
from datasets import find_dataset_def
from torch.utils.data import DataLoader
import numpy as np
import cv2
MVSDataset = find_dataset_def("dtu_yao")
dataset = MVSDataset("/home/xyguo/dataset/dtu_mvs/processed/mvs_training/dtu/", '../lists/dtu/train.txt', 'train',
3, 256)
dataloader = DataLoader(dataset, batch_size=2)
item = next(iter(dataloader))
imgs = item["imgs"][:, :, :, ::4, ::4].cuda()
proj_matrices = item["proj_matrices"].cuda()
mask = item["mask"].cuda()
depth = item["depth"].cuda()
depth_values = item["depth_values"].cuda()
imgs = torch.unbind(imgs, 1)
proj_matrices = torch.unbind(proj_matrices, 1)
ref_img, src_imgs = imgs[0], imgs[1:]
ref_proj, src_projs = proj_matrices[0], proj_matrices[1:]
if args.mode == "train":
if not os.path.isdir(args.logdir):
os.mkdir(args.logdir)
current_time_str = str(datetime.datetime.now().strftime('%Y%m%d_%H%M%S'))
print("current time", current_time_str)
print("creating new summary file")
logger = SummaryWriter(args.logdir)
print("argv:", sys.argv[1:])
print_args(args)
# dataset, dataloader
MVSDataset = find_dataset_def(args.dataset)
if args.dataset == 'dtu_yao':
train_dataset = MVSDataset(args.trainpath,
args.trainlist,
"train",
5,
robust_train=True)
test_dataset = MVSDataset(args.valpath,
args.vallist,
"val",
5,
robust_train=False)
TrainImgLoader = DataLoader(train_dataset,
args.batch_size,
shuffle=True,
if not os.path.exists(os.path.join(cfg.TEST.PATH, 'SyncedPoses.txt')):
logger.info("First run on this captured data, start the pre-processing...")
process_data(cfg.TEST.PATH)
else:
logger.info("Found SyncedPoses.txt, skipping data pre-processing...")
logger.info("Running NeuralRecon...")
transform = [transforms.ResizeImage((640, 480)),
transforms.ToTensor(),
transforms.RandomTransformSpace(
cfg.MODEL.N_VOX, cfg.MODEL.VOXEL_SIZE, random_rotation=False, random_translation=False,
paddingXY=0, paddingZ=0, max_epoch=cfg.TRAIN.EPOCHS),
transforms.IntrinsicsPoseToProjection(cfg.TEST.N_VIEWS, 4)]
transforms = transforms.Compose(transform)
ARKitDataset = find_dataset_def(cfg.DATASET)
test_dataset = ARKitDataset(cfg.TEST.PATH, "test", transforms, cfg.TEST.N_VIEWS, len(cfg.MODEL.THRESHOLDS) - 1)
data_loader = DataLoader(test_dataset, cfg.BATCH_SIZE, shuffle=False, num_workers=cfg.TEST.N_WORKERS, drop_last=False)
# model
logger.info("Initializing the model on GPU...")
model = NeuralRecon(cfg).cuda().eval()
model = torch.nn.DataParallel(model, device_ids=[0])
# use the latest checkpoint file
saved_models = [fn for fn in os.listdir(cfg.LOGDIR) if fn.endswith(".ckpt")]
saved_models = sorted(saved_models, key=lambda x: int(x.split('_')[-1].split('.')[0]))
loadckpt = os.path.join(cfg.LOGDIR, saved_models[-1])
logger.info("Resuming from " + str(loadckpt))
state_dict = torch.load(loadckpt)
model.load_state_dict(state_dict['model'], strict=False)
transforms.ResizeImage((640, 480)),
transforms.ToTensor(),
transforms.RandomTransformSpace(cfg.MODEL.N_VOX,
cfg.MODEL.VOXEL_SIZE,
random_rotation,
random_translation,
paddingXY,
paddingZ,
max_epoch=cfg.TRAIN.EPOCHS),
transforms.IntrinsicsPoseToProjection(n_views, 4),
]
transforms = transforms.Compose(transform)
# dataset, dataloader
MVSDataset = find_dataset_def(cfg.DATASET)
train_dataset = MVSDataset(cfg.TRAIN.PATH, "train", transforms,
cfg.TRAIN.N_VIEWS,
len(cfg.MODEL.THRESHOLDS) - 1)
test_dataset = MVSDataset(cfg.TEST.PATH, "test", transforms, cfg.TEST.N_VIEWS,
len(cfg.MODEL.THRESHOLDS) - 1)
if cfg.DISTRIBUTED:
train_sampler = DistributedSampler(train_dataset, shuffle=False)
TrainImgLoader = torch.utils.data.DataLoader(
train_dataset,
batch_size=cfg.BATCH_SIZE,
sampler=train_sampler,
num_workers=cfg.TRAIN.N_WORKERS,
pin_memory=True,
drop_last=True)
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 save_depth():
# dataset, dataloader
MVSDataset = find_dataset_def(args.dataset)
test_dataset = MVSDataset(args.testpath, args.testlist, "test",
args.numviews, 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))
# from collections import OrderedDict
# new_state_dict = OrderedDict()
state_dict = torch.load(args.loadckpt)
# for k, v in state_dict["model"].items():
# name = k.replace("module.", "") # remove module.
# new_state_dict[name] = v
# model.load_state_dict(new_state_dict) #, strict=False)
model.load_state_dict(state_dict['model']) #, strict=False)
model.eval()
with torch.no_grad():
for batch_idx, sample in enumerate(TestImgLoader):
orig_images = torch.unbind(sample["imgs"], 1)
ext_matrices = sample["ext_matrices"]
int_matrices = sample["int_matrices"]
ref_img = orig_images[0].squeeze().numpy()
n_img = orig_images[1].squeeze().numpy()
ref_ext = ext_matrices[0].squeeze().numpy()
ref_int = int_matrices[0].squeeze().numpy()
n_ext = ext_matrices[1].squeeze().numpy()
n_int = int_matrices[1].squeeze().numpy()
sample_cuda = tocuda(sample)
outputs = model(sample_cuda["imgs"], sample_cuda["proj_matrices"],
sample_cuda["depth_values"])
outputs = tensor2numpy(outputs)
del sample_cuda
ref_depth = cv2.resize(outputs["depth"].squeeze(),
(ref_img.shape[2], ref_img.shape[1]))
x = np.linspace(0, ref_img.shape[2], ref_img.shape[2], axis=-1)
y = np.linspace(0, ref_img.shape[1], ref_img.shape[1], axis=-1)
xv, yv = np.meshgrid(x, y)
W = PixelCoordToWorldCoord(ref_int, ref_ext[:3, :3],
ref_ext[:3, 3:4], xv, yv, ref_depth)
xp, yp = WorldCoordTopixelCoord(n_int, n_ext[:3, :3],
n_ext[:3, 3:4], W)
# xp, yp = WorldCoordTopixelCoord(ref_int, ref_ext[:3,:3], ref_ext[:3, 3:4], W)
projected_img = GetImageAtPixelCoordinates(ref_img, xp, yp)
fig = plt.figure()
plt.subplot(1, 3, 1)
plt.imshow(np.swapaxes(ref_img, 0, 2).swapaxes(0, 1))
plt.subplot(1, 3, 2)
plt.imshow(np.swapaxes(n_img, 0, 2).swapaxes(0, 1))
plt.subplot(1, 3, 3)
plt.imshow(np.swapaxes(projected_img, 0, 2).swapaxes(0, 1))
plt.show()
plt.tight_layout()
print('Iter {}/{}'.format(batch_idx, len(TestImgLoader)))
filenames = sample["filename"]
