def test(model_path):
test_args = TestOptions().parse()
test_args.thread = 0
test_args.batchsize = 1
merge_cfg_from_file(test_args)
data_loader = CustomerDataLoader(test_args)
test_datasize = len(data_loader)
logger.info('{:>15}: {:<30}'.format('test_data_size', test_datasize))
# load model
model = MetricDepthModel()
model.eval()
test_args.load_ckpt = model_path
# load checkpoint
if test_args.load_ckpt:
load_ckpt(test_args, model)
model.cuda()
# model = torch.nn.DataParallel(model)
# test
smoothed_absRel = SmoothedValue(test_datasize)
smoothed_rms = SmoothedValue(test_datasize)
smoothed_logRms = SmoothedValue(test_datasize)
smoothed_squaRel = SmoothedValue(test_datasize)
smoothed_silog = SmoothedValue(test_datasize)
smoothed_silog2 = SmoothedValue(test_datasize)
smoothed_log10 = SmoothedValue(test_datasize)
smoothed_delta1 = SmoothedValue(test_datasize)
smoothed_delta2 = SmoothedValue(test_datasize)
smoothed_delta3 = SmoothedValue(test_datasize)
smoothed_whdr = SmoothedValue(test_datasize)
smoothed_criteria = {
'err_absRel': smoothed_absRel,
'err_squaRel': smoothed_squaRel,
'err_rms': smoothed_rms,
'err_silog': smoothed_silog,
'err_logRms': smoothed_logRms,
'err_silog2': smoothed_silog2,
'err_delta1': smoothed_delta1,
'err_delta2': smoothed_delta2,
'err_delta3': smoothed_delta3,
'err_log10': smoothed_log10,
'err_whdr': smoothed_whdr
}
for i, data in enumerate(data_loader):
out = model.inference(data)
pred_depth = torch.squeeze(out['b_fake'])
img_path = data['A_paths']
invalid_side = data['invalid_side'][0]
pred_depth = pred_depth[invalid_side[0]:pred_depth.size(0) -
invalid_side[1], :]
pred_depth = pred_depth / data['ratio'].cuda() # scale the depth
pred_depth = resize_image(pred_depth,
torch.squeeze(data['B_raw']).shape)
smoothed_criteria = evaluate_err(pred_depth,
data['B_raw'],
smoothed_criteria,
mask=(45, 471, 41, 601),
scale=10.)
# save images
model_name = test_args.load_ckpt.split('/')[-1].split('.')[0]
image_dir = os.path.join(cfg.ROOT_DIR, './evaluation',
cfg.MODEL.ENCODER, model_name)
if not os.path.exists(image_dir):
os.makedirs(image_dir)
img_name = img_path[0].split('/')[-1]
#plt.imsave(os.path.join(image_dir, 'd_' + img_name), pred_depth, cmap='rainbow')
#cv2.imwrite(os.path.join(image_dir, 'rgb_' + img_name), data['A_raw'].numpy().squeeze())
# print('processing (%04d)-th image... %s' % (i, img_path))
# print("###############absREL ERROR: %f", smoothed_criteria['err_absRel'].GetGlobalAverageValue())
# print("###############silog ERROR: %f", np.sqrt(smoothed_criteria['err_silog2'].GetGlobalAverageValue() - (
# smoothed_criteria['err_silog'].GetGlobalAverageValue()) ** 2))
# print("###############log10 ERROR: %f", smoothed_criteria['err_log10'].GetGlobalAverageValue())
# print("###############RMS ERROR: %f", np.sqrt(smoothed_criteria['err_rms'].GetGlobalAverageValue()))
# print("###############delta_1 ERROR: %f", smoothed_criteria['err_delta1'].GetGlobalAverageValue())
# print("###############delta_2 ERROR: %f", smoothed_criteria['err_delta2'].GetGlobalAverageValue())
# print("###############delta_3 ERROR: %f", smoothed_criteria['err_delta3'].GetGlobalAverageValue())
# print("###############squaRel ERROR: %f", smoothed_criteria['err_squaRel'].GetGlobalAverageValue())
# print("###############logRms ERROR: %f", np.sqrt(smoothed_criteria['err_logRms'].GetGlobalAverageValue()))
f.write("tested model:" + model_path)
f.write('\n')
f.write("###############absREL ERROR:" +
str(smoothed_criteria['err_absRel'].GetGlobalAverageValue()))
f.write('\n')
f.write("###############silog ERROR:" + str(
np.sqrt(smoothed_criteria['err_silog2'].GetGlobalAverageValue() -
(smoothed_criteria['err_silog'].GetGlobalAverageValue())**2)))
f.write('\n')
f.write("###############log10 ERROR:" +
str(smoothed_criteria['err_log10'].GetGlobalAverageValue()))
f.write('\n')
f.write("###############RMS ERROR:" +
str(np.sqrt(smoothed_criteria['err_rms'].GetGlobalAverageValue())))
f.write('\n')
f.write("###############delta_1 ERROR:" +
str(smoothed_criteria['err_delta1'].GetGlobalAverageValue()))
f.write('\n')
f.write("###############delta_2 ERROR:" +
str(smoothed_criteria['err_delta2'].GetGlobalAverageValue()))
f.write('\n')
f.write("###############delta_3 ERROR:" +
str(smoothed_criteria['err_delta3'].GetGlobalAverageValue()))
f.write('\n')
f.write("###############squaRel ERROR:" +
str(smoothed_criteria['err_squaRel'].GetGlobalAverageValue()))
f.write('\n')
f.write(
"###############logRms ERROR:" +
str(np.sqrt(smoothed_criteria['err_logRms'].GetGlobalAverageValue())))
f.write('\n')
f.write(
'-----------------------------------------------------------------------------'
)
f.write('\n')
loss_func = ModelLoss()
val_err = [{'abs_rel': 0, 'silog': 0}]
ignore_step = -1
# Lerning strategy
lr_optim_lambda = lambda iter: (1.0 - iter / (float(total_iters)))**0.9
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer.optimizer,
lr_lambda=lr_optim_lambda)
#scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer.optimizer, milestones=train_args.epoch, gamma=0.1)
# load checkpoint
if train_args.load_ckpt:
load_ckpt(train_args, model, optimizer.optimizer, scheduler, val_err)
ignore_step = train_args.start_step - train_args.start_epoch * math.ceil(
train_datasize / train_args.batchsize)
if gpu_num != -1:
model = torch.nn.DataParallel(model)
try:
for epoch in range(train_args.start_epoch, cfg.TRAIN.EPOCH[-1]):
# training
train(train_dataloader, model, epoch, loss_func, optimizer,
scheduler, training_stats, val_dataloader, val_err,
ignore_step)
ignore_step = -1
except (RuntimeError, KeyboardInterrupt):
logger.info('Save ckpt on exception ...')
test_args = TestOptions().parse()
test_args.thread = 1
test_args.batchsize = 1
merge_cfg_from_file(test_args)
data_loader = CustomerDataLoader(test_args)
test_datasize = len(data_loader)
logger.info('{:>15}: {:<30}'.format('test_data_size', test_datasize))
# load model
model = MetricDepthModel()
model.eval()
# load checkpoint
if test_args.load_ckpt:
load_ckpt(test_args, model)
model.cuda()
model = torch.nn.DataParallel(model)
# test
smoothed_absRel = SmoothedValue(test_datasize)
smoothed_rms = SmoothedValue(test_datasize)
smoothed_logRms = SmoothedValue(test_datasize)
smoothed_squaRel = SmoothedValue(test_datasize)
smoothed_silog = SmoothedValue(test_datasize)
smoothed_silog2 = SmoothedValue(test_datasize)
smoothed_log10 = SmoothedValue(test_datasize)
smoothed_delta1 = SmoothedValue(test_datasize)
smoothed_delta2 = SmoothedValue(test_datasize)
smoothed_delta3 = SmoothedValue(test_datasize)
smoothed_whdr = SmoothedValue(test_datasize)
def main():
test_args = TestOptions().parse()
test_args.thread = 1 # test code only supports thread = 1
test_args.batchsize = 1 # test code only supports batchSize = 1
data_loader = CustomerDataLoader(test_args)
test_datasize = len(data_loader)
logger.info('{:>15}: {:<30}'.format('test_data_size', test_datasize))
# load model
model = DepthNormal()
# evaluate mode
model.eval()
# load checkpoint
if test_args.load_ckpt:
load_ckpt(test_args, model)
model.cuda()
model = torch.nn.DataParallel(model)
for i, data in enumerate(data_loader):
out = model.module.inference(data)
pred_depth = np.squeeze(out['b_fake']) * 80. # [h, w]
pred_conf = np.squeeze(out['b_fake_conf']) # [c, h, w]
# the image size has been padded to the size (385, 1243)
pred_depth_crop = pred_depth[data['pad_raw'][0][0]:,
data['pad_raw'][0][2]:]
pred_conf_crop = pred_conf[:, data['pad_raw'][0][0]:,
data['pad_raw'][0][2]:]
sample_th = 0.15
sample_mask = get_sample_mask(pred_conf_crop.cpu().numpy(),
threshold=sample_th) # [h, w]
#######################################################################################
# add by users
img_name = data['A_paths'][0].split('/')[-1][:-4]
calib_name = img_name + '.txt'
calib_dir = os.path.join(calib_fold, calib_name)
camera_para = np.genfromtxt(calib_dir,
delimiter=' ',
skip_footer=3,
dtype=None)
P3_0 = camera_para[3]
P2_0 = camera_para[2]
P3_2 = P3_0
P3_2[4] -= P2_0[4]
R0_rect = np.genfromtxt(calib_dir,
delimiter=' ',
skip_header=4,
skip_footer=2)
Tr_velo_to_cam0 = np.genfromtxt(calib_dir,
delimiter=' ',
skip_header=5,
skip_footer=1)
pcd_cam2 = reconstruct_3D(pred_depth_crop.cpu().numpy(), P3_2[3],
P3_2[7], P3_2[1], P3_2[6])
# Transfer points in cam2 coordinate to cam0 coordinate
pcd_cam0 = pcd_cam2 - np.array([[[P2_0[4] / P2_0[1]]],
[[P2_0[8] / P2_0[1]]],
[[P2_0[12] / P2_0[1]]]])
# Transfer points in cam0 coordinate to velo coordinate
pcd_velo = transfer_points_in_cam0_to_velo(pcd_cam0, R0_rect,
Tr_velo_to_cam0)
rgb = data['A_raw'][0].cpu().numpy()
save_ply(pcd_velo,
rgb,
os.path.join(pcd_folder, img_name) + '_sample.ply',
sample_mask=sample_mask)
#save_ply(pcd_cam2, rgb, os.path.join(pcd_folder, img_name) + '.ply')
print('saved', img_name)
else:
img = img.astype(np.float32)
img = torch.from_numpy(img)
return img
if __name__ == '__main__':
args = parse_args()
# create depth model
depth_model = RelDepthModel()
depth_model.eval()
# load checkpoint
load_ckpt(args, depth_model, None, None)
depth_model.cuda()
image_dir = os.path.dirname(os.path.dirname(__file__)) + '/test_images/'
imgs_list = os.listdir(image_dir)
imgs_list.sort()
imgs_path = [
os.path.join(image_dir, i) for i in imgs_list if i != 'outputs'
]
image_dir_out = image_dir + '/outputs'
os.makedirs(image_dir_out, exist_ok=True)
for i, v in enumerate(imgs_path):
print('processing (%04d)-th image... %s' % (i, v))
rgb = cv2.imread(v)
rgb_c = rgb[:, :, ::-1].copy()
def main_worker(local_rank: int, ngpus_per_node: int, train_args, val_args):
train_args.global_rank = train_args.node_rank * ngpus_per_node + local_rank
train_args.local_rank = local_rank
val_args.global_rank = train_args.global_rank
val_args.local_rank = local_rank
merge_cfg_from_file(train_args)
global logger
# Set logger
log_output_dir = cfg.TRAIN.LOG_DIR
if log_output_dir:
try:
os.makedirs(log_output_dir)
except OSError as e:
if e.errno != errno.EEXIST:
raise
logger = setup_distributed_logger("lib", log_output_dir, local_rank,
cfg.TRAIN.RUN_NAME + '.txt')
tblogger = None
if train_args.use_tfboard and local_rank == 0:
from tensorboardX import SummaryWriter
tblogger = SummaryWriter(cfg.TRAIN.LOG_DIR)
# init
if train_args.distributed:
torch.cuda.set_device(local_rank)
dist.init_process_group(backend='nccl',
init_method=train_args.dist_url,
world_size=train_args.world_size,
rank=train_args.global_rank)
# load model
model = RelDepthModel()
if train_args.world_size > 1:
assert is_pytorch_1_1_0_or_later(), \
"SyncBatchNorm is only available in pytorch >= 1.1.0"
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
logger.info('Using SyncBN!')
if train_args.distributed:
model = torch.nn.parallel.DistributedDataParallel(
model.cuda(), device_ids=[local_rank], output_device=local_rank)
else:
model = torch.nn.DataParallel(model.cuda())
val_err = [{'abs_rel': 0, 'whdr': 0}]
# Print configs and logs
print_options(train_args, logger)
# training and validation dataloader
train_dataloader, train_sample_size = MultipleDataLoaderDistributed(
train_args)
val_dataloader, val_sample_size = MultipleDataLoaderDistributed(val_args)
cfg.TRAIN.LR_SCHEDULER_MULTISTEPS = np.array(
train_args.lr_scheduler_multiepochs) * math.ceil(
np.sum(train_sample_size) /
(train_args.world_size * train_args.batchsize))
# Optimizer
optimizer = ModelOptimizer(model)
# lr_optim_lambda = lambda iter: (1.0 - iter / (float(total_iters))) ** 0.9
# scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer.optimizer, lr_lambda=lr_optim_lambda)
# lr_scheduler_step = 15000
# scheduler = torch.optim.lr_scheduler.StepLR(optimizer.optimizer, step_size=lr_scheduler_step, gamma=0.9)
scheduler = make_lr_scheduler(cfg=cfg, optimizer=optimizer.optimizer)
if train_args.load_ckpt:
load_ckpt(train_args, model, optimizer.optimizer, scheduler, val_err)
# obtain the current sample ratio
sample_ratio = increase_sample_ratio_steps(
train_args.start_step,
base_ratio=train_args.sample_start_ratio,
step_size=train_args.sample_ratio_steps)
# reconstruct the train_dataloader with the new sample_ratio
train_dataloader, train_sample_size = MultipleDataLoaderDistributed(
train_args, sample_ratio=sample_ratio)
total_iters = math.ceil(
np.sum(train_sample_size) /
(train_args.world_size * train_args.batchsize)) * train_args.epoch
cfg.TRAIN.MAX_ITER = total_iters
cfg.TRAIN.GPU_NUM = train_args.world_size
print_configs(cfg)
save_to_disk = main_process(train_args.distributed, local_rank)
do_train(train_dataloader, val_dataloader, train_args, model, save_to_disk,
scheduler, optimizer, val_err, logger, tblogger)
def __init__(self):
self.device = "cuda" if torch.cuda.is_available() else "cpu"
# Setup AdaBins model
self.adabins_nyu_infer_helper = InferenceHelper(dataset='nyu',
device=self.device)
self.adabins_kitti_infer_helper = InferenceHelper(dataset='kitti',
device=self.device)
# Setup DiverseDepth model
class DiverseDepthArgs:
def __init__(self):
self.resume = False
self.cfg_file = "lib/configs/resnext50_32x4d_diversedepth_regression_vircam"
self.load_ckpt = "pretrained/DiverseDepth.pth"
diverse_depth_args = DiverseDepthArgs()
merge_cfg_from_file(diverse_depth_args)
self.diverse_depth_model = RelDepthModel()
self.diverse_depth_model.eval()
# load checkpoint
load_ckpt(diverse_depth_args, self.diverse_depth_model)
# TODO: update this - see how `device` argument should be processsed
if self.device == "cuda":
self.diverse_depth_model.cuda()
self.diverse_depth_model = torch.nn.DataParallel(
self.diverse_depth_model)
# Setup MiDaS model
self.midas_model_path = "./pretrained/MiDaS_f6b98070.pt"
midas_model_type = "large"
# load network
if midas_model_type == "large":
self.midas_model = MidasNet(self.midas_model_path,
non_negative=True)
self.midas_net_w, self.midas_net_h = 384, 384
elif midas_model_type == "small":
self.midas_model = MidasNet_small(self.midas_model_path,
features=64,
backbone="efficientnet_lite3",
exportable=True,
non_negative=True,
blocks={'expand': True})
self.midas_net_w, self.midas_net_h = 256, 256
self.midas_transform = Compose([
Resize(
self.midas_net_w,
self.midas_net_h,
resize_target=None,
keep_aspect_ratio=True,
ensure_multiple_of=32,
resize_method="upper_bound",
image_interpolation_method=cv2.INTER_CUBIC,
),
NormalizeImage(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
PrepareForNet(),
])
self.midas_model.eval()
self.midas_optimize = True
if self.midas_optimize == True:
rand_example = torch.rand(1, 3, self.midas_net_h, self.midas_net_w)
self.midas_model(rand_example)
traced_script_module = torch.jit.trace(self.midas_model,
rand_example)
self.midas_model = traced_script_module
if self.device == "cuda":
self.midas_model = self.midas_model.to(
memory_format=torch.channels_last)
self.midas_model = self.midas_model.half()
self.midas_model.to(torch.device(self.device))
# Setup SGDepth model
self.sgdepth_model = InferenceEngine.SgDepthInference()
# Setup monodepth2 model
self.monodepth2_model_path = "pretrained/monodepth2_mono+stereo_640x192"
monodepth2_device = torch.device(self.device)
encoder_path = os.path.join(self.monodepth2_model_path, "encoder.pth")
depth_decoder_path = os.path.join(self.monodepth2_model_path,
"depth.pth")
# LOADING PRETRAINED MODEL
print(" Loading Monodepth2 pretrained encoder")
self.monodepth2_encoder = networks.ResnetEncoder(18, False)
loaded_dict_enc = torch.load(encoder_path,
map_location=monodepth2_device)
# extract the height and width of image that this model was trained with
self.feed_height = loaded_dict_enc['height']
self.feed_width = loaded_dict_enc['width']
filtered_dict_enc = {
k: v
for k, v in loaded_dict_enc.items()
if k in self.monodepth2_encoder.state_dict()
}
self.monodepth2_encoder.load_state_dict(filtered_dict_enc)
self.monodepth2_encoder.to(monodepth2_device)
self.monodepth2_encoder.eval()
print(" Loading pretrained decoder")
self.monodepth2_depth_decoder = networks.DepthDecoder(
num_ch_enc=self.monodepth2_encoder.num_ch_enc, scales=range(4))
loaded_dict = torch.load(depth_decoder_path,
map_location=monodepth2_device)
self.monodepth2_depth_decoder.load_state_dict(loaded_dict)
self.monodepth2_depth_decoder.to(monodepth2_device)
self.monodepth2_depth_decoder.eval()
def train(local_rank, distributed, train_args, logger, tblogger=None):
# load model
model = RelDepthModel()
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
if cfg.MODEL.USE_SYNCBN:
assert is_pytorch_1_1_0_or_later(), \
"SyncBatchNorm is only available in pytorch >= 1.1.0"
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
# Optimizer
optimizer = ModelOptimizer(model)
#lr_optim_lambda = lambda iter: (1.0 - iter / (float(total_iters))) ** 0.9
#scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer.optimizer, lr_lambda=lr_optim_lambda)
lr_scheduler_step = 15000
scheduler = torch.optim.lr_scheduler.StepLR(optimizer.optimizer,
step_size=lr_scheduler_step,
gamma=0.9)
if distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[local_rank], output_device=local_rank)
val_err = [{'abs_rel': 0, 'whdr': 0}]
# training and validation dataloader
val_dataloader = MultipleDataLoaderDistributed(val_args)
if train_args.load_ckpt:
load_ckpt(train_args, model, optimizer.optimizer, scheduler, val_err)
# obtain the current sample ratio
sample_ratio = increase_sample_ratio_steps(
train_args.start_step,
base_ratio=train_args.sample_start_ratio,
step_size=train_args.sample_ratio_steps)
# reconstruct the train_dataloader with the new sample_ratio
train_dataloader = MultipleDataLoaderDistributed(
train_args, sample_ratio=sample_ratio)
if not train_args.resume:
scheduler.__setattr__('step_size', lr_scheduler_step)
else:
train_dataloader = MultipleDataLoaderDistributed(train_args)
train_datasize = len(train_dataloader)
val_datasize = len(val_dataloader)
merge_cfg_from_file(train_args)
# total iterations
total_iters = math.ceil(
train_datasize / train_args.batchsize) * train_args.epoch
cfg.TRAIN.MAX_ITER = total_iters
cfg.TRAIN.GPU_NUM = gpu_num
# Print configs and logs
if get_rank() == 0:
train_opt.print_options(train_args)
val_opt.print_options(val_args)
print_configs(cfg)
logger.info('{:>15}: {:<30}'.format('GPU_num', gpu_num))
logger.info('{:>15}: {:<30}'.format('train_data_size', train_datasize))
logger.info('{:>15}: {:<30}'.format('val_data_size', val_datasize))
logger.info('{:>15}: {:<30}'.format('total_iterations', total_iters))
save_to_disk = get_rank() == 0
do_train(train_dataloader, val_dataloader, train_args, model, save_to_disk,
scheduler, optimizer, val_err, logger, tblogger)
