def init_model(args):
print("init model start...................")
hvd_disable()
# Parse arguments
config, state_dict = parse_test_file(args.checkpoint)
# If no image shape is provided, use the checkpoint one
image_shape = args.image_shape
if image_shape is None:
image_shape = config.datasets.augmentation.image_shape
print(f"input image shape:{image_shape}")
# Set debug if requested
set_debug(config.debug)
# Initialize model wrapper from checkpoint arguments
model_wrapper = ModelWrapper(config, load_datasets=False)
# Restore monodepth_model state
model_wrapper.load_state_dict(state_dict)
# Send model to GPU if available
if torch.cuda.is_available():
model_wrapper = model_wrapper.to('cuda')
else:
raise RuntimeError("cuda is not available")
# Set to eval mode
model_wrapper.eval()
print("init model finish...................")
return model_wrapper, image_shape
def test(ckpt_file, cfg_file, half):
"""
Monocular depth estimation test script.
Parameters
----------
ckpt_file : str
Checkpoint path for a pretrained model
cfg_file : str
Configuration file
half: bool
use half precision (fp16)
"""
# Initialize horovod
hvd_init()
# Parse arguments
config, state_dict = parse_test_file(ckpt_file, cfg_file)
# Set debug if requested
set_debug(config.debug)
# Initialize monodepth model from checkpoint arguments
model_wrapper = ModelWrapper(config)
# Restore model state
model_wrapper.load_state_dict(state_dict)
# change to half precision for evaluation if requested
config.arch["dtype"] = torch.float16 if half else None
# Create trainer with args.arch parameters
trainer = HorovodTrainer(**config.arch)
# Test model
trainer.test(model_wrapper)
def main(args):
# Initialize horovod
hvd_init()
# Parse arguments
config, state_dict = parse_test_file(args.checkpoint)
# If no image shape is provided, use the checkpoint one
image_shape = args.image_shape
if image_shape is None:
image_shape = config.datasets.augmentation.image_shape
# Set debug if requested
set_debug(config.debug)
# Initialize model wrapper from checkpoint arguments
model_wrapper = ModelWrapper(config, load_datasets=False)
# Restore monodepth_model state
model_wrapper.load_state_dict(state_dict)
# change to half precision for evaluation if requested
dtype = torch.float16 if args.half else None
# Send model to GPU if available
if torch.cuda.is_available():
model_wrapper = model_wrapper.to('cuda:{}'.format(rank()), dtype=dtype)
# Set to eval mode
model_wrapper.eval()
if os.path.isdir(args.input):
# If input file is a folder, search for image files
files = []
for ext in ['png', 'jpg']:
files.extend(glob((os.path.join(args.input, '*.{}'.format(ext)))))
files.sort()
print0('Found {} files'.format(len(files)))
else:
# Otherwise, use it as is
files = [args.input]
# Process each file
for fn in files[rank()::world_size()]:
infer_and_save_depth(fn, args.output, model_wrapper, image_shape,
args.half, args.save)
def train(file, config):
"""
Monocular depth estimation training script.
Parameters
----------
file : str
Filepath, can be either a
**.yaml** for a yacs configuration file or a
**.ckpt** for a pre-trained checkpoint file.
"""
# Initialize horovod
hvd_init()
# Produce configuration and checkpoint from filename
config, ckpt = parse_train_file(file, config)
# Set debug if requested
set_debug(config.debug)
# Wandb Logger
logger = None if config.wandb.dry_run or rank() > 0 \
else filter_args_create(WandbLogger, config.wandb)
# model checkpoint
checkpoint = None if config.checkpoint.filepath is '' or rank() > 0 else \
filter_args_create(ModelCheckpoint, config.checkpoint)
# Initialize model wrapper
model_wrapper = ModelWrapper(config, resume=ckpt, logger=logger)
# Create trainer with args.arch parameters
trainer = HorovodTrainer(**config.arch, checkpoint=checkpoint)
# Train model
trainer.fit(model_wrapper)
def main(args):
# Initialize horovod
hvd_init()
# Parse arguments
config, state_dict = parse_test_file(args.checkpoint)
# If no image shape is provided, use the checkpoint one
image_shape = args.image_shape
if image_shape is None:
image_shape = config.datasets.augmentation.image_shape
print(image_shape)
# Set debug if requested
set_debug(config.debug)
# Initialize model wrapper from checkpoint arguments
model_wrapper = ModelWrapper(config, load_datasets=False)
# Restore monodepth_model state
model_wrapper.load_state_dict(state_dict)
# change to half precision for evaluation if requested
dtype = torch.float16 if args.half else None
# Send model to GPU if available
if torch.cuda.is_available():
model_wrapper = model_wrapper.to('cuda:{}'.format(rank()), dtype=dtype)
# Set to eval mode
model_wrapper.eval()
if os.path.isdir(args.input):
# If input file is a folder, search for image files
files = []
for ext in ['png', 'jpg']:
files.extend(glob((os.path.join(args.input, '*.{}'.format(ext)))))
files.sort()
print0('Found {} files'.format(len(files)))
else:
raise RuntimeError("Input needs directory, not file")
if not os.path.isdir(args.output):
root, file_name = os.path.split(args.output)
os.makedirs(root, exist_ok=True)
else:
raise RuntimeError("Output needs to be a file")
# Process each file
list_of_files = list(
zip(files[rank():-2:world_size()], files[rank() + 1:-1:world_size()],
files[rank() + 2::world_size()]))
if args.offset:
list_of_files = list_of_files[args.offset:]
if args.limit:
list_of_files = list_of_files[:args.limit]
for fn1, fn2, fn3 in list_of_files:
infer_and_save_pose([fn1, fn3], fn2, model_wrapper, image_shape,
args.half, args.save)
position = np.zeros(3)
orientation = np.eye(3)
for key in sorted(poses.keys()):
rot_matrix, translation = poses[key]
orientation = orientation.dot(rot_matrix.tolist())
position += orientation.dot(translation.tolist())
poses[key] = {
"rot":
rot_matrix.tolist(),
"trans":
translation.tolist(),
"pose": [
*orientation[0], position[0], *orientation[1], position[1],
*orientation[2], position[2], 0, 0, 0, 1
]
}
json.dump(poses, open(args.output, "w"), sort_keys=True)
print(f"Written pose of {len(list_of_files)} images to {args.output}")