from tools.parse_arg_test import TestOptions
from lib.core.config import merge_cfg_from_file
from data.load_dataset import CustomerDataLoader
from lib.models.image_transfer import resize_image
from lib.utils.evaluate_depth_error import evaluate_err
from lib.models.metric_depth_model import MetricDepthModel
from lib.utils.logging import setup_logging, SmoothedValue
import matplotlib.pyplot as plt
logger = setup_logging(__name__)
if __name__ == '__main__':
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)
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')
bg_smoothed_criteria['err_absRel'].GetGlobalAverageValue(),
'silog':
np.sqrt(bg_smoothed_criteria['err_silog2'].GetGlobalAverageValue() -
(bg_smoothed_criteria['err_silog'].GetGlobalAverageValue())**2)
}
print("global: ", val_metrics)
print("roi: ", rois_val_metrics)
print("bg: ", bg_val_metrics)
return val_metrics
if __name__ == '__main__':
train_dataloader = CustomerDataLoader(train_args)
train_datasize = len(train_dataloader)
gpu_num = torch.cuda.device_count()
merge_cfg_from_file(train_datasize, gpu_num)
val_dataloader = CustomerDataLoader(val_args)
val_datasize = len(val_dataloader)
# tensorboard logger
if train_args.use_tfboard:
from tensorboardX import SummaryWriter
tblogger = SummaryWriter(cfg.TRAIN.LOG_DIR)
# training status for logging
training_stats = TrainingStats(
train_args, cfg.TRAIN.LOG_INTERVAL,
tblogger if train_args.use_tfboard else None)
# total iterations
pred_depth = pred_depth[invalid_side[0]:pred_depth.size(0) - invalid_side[1],
invalid_side[2]:pred_depth.size(1) - invalid_side[3]]
pred_depth_resize = resize_image(pred_depth, torch.squeeze(data['B_raw']).shape)
pred_depth_metric = recover_metric_depth(pred_depth_resize, data['B_raw'])
smoothed_criteria = validate_rel_depth_err(pred_depth_metric, data['B_raw'], smoothed_criteria, scale=1.0)
return {'abs_rel': smoothed_criteria['err_absRel'].GetGlobalAverageValue(),
'whdr': smoothed_criteria['err_whdr'].GetGlobalAverageValue()}
if __name__=='__main__':
# Train args
train_opt = TrainOptions()
train_args = train_opt.parse()
merge_cfg_from_file(train_args)
gpu_num = torch.cuda.device_count()
cfg.TRAIN.GPU_NUM = gpu_num
# Validation args
val_opt = ValOptions()
val_args = val_opt.parse()
val_args.batchsize = 1
val_args.thread = 0
# Logger
log_output_dir = cfg.TRAIN.LOG_DIR
if log_output_dir:
try:
os.makedirs(log_output_dir)
assert bool(args.load_ckpt) ^ bool(args.load_detectron), \
'Exactly one of --load_ckpt and --load_detectron should be specified.'
if args.output_dir is None:
ckpt_path = args.load_ckpt if args.load_ckpt else args.load_detectron
args.output_dir = os.path.join(
os.path.dirname(os.path.dirname(ckpt_path)), 'test')
logger.info('Automatically set output directory to %s',
args.output_dir)
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
cfg.VIS = args.vis
if args.cfg_file is not None:
merge_cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
merge_cfg_from_list(args.set_cfgs)
if args.dataset == "coco2017":
cfg.TEST.DATASETS = ('coco_2017_val', )
cfg.MODEL.NUM_CLASSES = 81
elif args.dataset == "keypoints_coco2017":
cfg.TEST.DATASETS = ('keypoints_coco_2017_val', )
cfg.MODEL.NUM_CLASSES = 2
else: # For subprocess call
assert cfg.TEST.DATASETS, 'cfg.TEST.DATASETS shouldn\'t be empty'
assert_and_infer_cfg()
logger.info('Testing with config:')
logger.info(pprint.pformat(cfg))
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)
