def setup_percep_net(config, prepared, **kwargs):
"""
Create a perceputal network
Parameters
----------
config : CfgNode
Network configuration
prepared : bool
True if the network has been prepared before
kwargs : dict
Extra parameters for the network
Returns
-------
depth_net : nn.Module
Create depth network
"""
print0(pcolor('PercepNet: %s' % config.name, 'yellow'))
percep_net = load_class_args_create(
config.name,
paths=[
'dro_sfm.networks.layers',
],
args={
**config,
**kwargs
},
)
return percep_net
def setup_pose_net(config, prepared, **kwargs):
"""
Create a pose network
Parameters
----------
config : CfgNode
Network configuration
prepared : bool
True if the network has been prepared before
kwargs : dict
Extra parameters for the network
Returns
-------
pose_net : nn.Module
Created pose network
"""
print0(pcolor('PoseNet: %s' % config.name, 'yellow'))
pose_net = load_class_args_create(
config.name,
paths=[
'dro_sfm.networks.pose',
],
args={
**config,
**kwargs
},
)
if not prepared and config.checkpoint_path is not '':
pose_net = load_network(pose_net, config.checkpoint_path,
['pose_net', 'pose_network'])
return pose_net
def load_network(network, path, prefixes=''):
"""
Loads a pretrained network
Parameters
----------
network : nn.Module
Network that will receive the pretrained weights
path : str
File containing a 'state_dict' key with pretrained network weights
prefixes : str or list of str
Layer name prefixes to consider when loading the network
Returns
-------
network : nn.Module
Updated network with pretrained weights
"""
prefixes = make_list(prefixes)
# If path is a string
if is_str(path):
saved_state_dict = torch.load(path, map_location='cpu')['state_dict']
if path.endswith('.pth.tar'):
saved_state_dict = backwards_state_dict(saved_state_dict)
# If state dict is already provided
else:
saved_state_dict = path
# Get network state dict
network_state_dict = network.state_dict()
updated_state_dict = OrderedDict()
n, n_total = 0, len(network_state_dict.keys())
for key, val in saved_state_dict.items():
for prefix in prefixes:
prefix = prefix + '.'
if prefix in key:
idx = key.find(prefix) + len(prefix)
key = key[idx:]
if key in network_state_dict.keys() and \
same_shape(val.shape, network_state_dict[key].shape):
updated_state_dict[key] = val
n += 1
try:
network.load_state_dict(updated_state_dict, strict=True)
except Exception as e:
print(e)
network.load_state_dict(updated_state_dict, strict=False)
base_color, attrs = 'cyan', ['bold', 'dark']
color = 'green' if n == n_total else 'yellow' if n > 0 else 'red'
print0(pcolor('=====###### Pretrained {} loaded:'.format(prefixes[0]), base_color, attrs=attrs) +
pcolor(' {}/{} '.format(n, n_total), color, attrs=attrs) +
pcolor('tensors', base_color, attrs=attrs))
return network
def prepare_model(self, resume=None):
"""Prepare self.model (incl. loading previous state)"""
print0(pcolor('### Preparing Model', 'green'))
self.model = setup_model(self.config.model, self.config.prepared)
# Resume model if available
if resume:
print0(
pcolor('### Resuming from {}'.format(resume['file']),
'magenta',
attrs=['bold']))
self.model = load_network(self.model, resume['state_dict'],
'model')
if 'epoch' in resume:
self.current_epoch = resume['epoch']
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 prepare_datasets(self, validation_requirements, test_requirements):
"""Prepare datasets for training, validation and test."""
# Prepare datasets
print0(pcolor('### Preparing Datasets', 'green'))
augmentation = self.config.datasets.augmentation
# Setup train dataset (requirements are given by the model itself)
self.train_dataset = setup_dataset(self.config.datasets.train, 'train',
self.model.train_requirements,
**augmentation)
# Setup validation dataset
self.validation_dataset = setup_dataset(
self.config.datasets.validation, 'validation',
validation_requirements, **augmentation)
# Setup test dataset
self.test_dataset = setup_dataset(self.config.datasets.test, 'test',
test_requirements, **augmentation)
def setup_model(config, prepared, **kwargs):
"""
Create a model
Parameters
----------
config : CfgNode
Model configuration (cf. configs/default_config.py)
prepared : bool
True if the model has been prepared before
kwargs : dict
Extra parameters for the model
Returns
-------
model : nn.Module
Created model
"""
print0(pcolor('Model: %s' % config.name, 'yellow'))
config.loss.min_depth = config.params.min_depth
config.loss.max_depth = config.params.max_depth
model = load_class(config.name, paths=[
'dro_sfm.models',
])(**{
**config.loss,
**kwargs
})
# Add depth network if required
if model.network_requirements['depth_net']:
config.depth_net.max_depth = config.params.max_depth
config.depth_net.min_depth = config.params.min_depth
model.add_depth_net(setup_depth_net(config.depth_net, prepared))
# Add pose network if required
if model.network_requirements['pose_net']:
model.add_pose_net(setup_pose_net(config.pose_net, prepared))
# Add percep_net if required
if model.network_requirements['percep_net']:
model.add_percep_net(setup_percep_net(config.percep_net, prepared))
# If a checkpoint is provided, load pretrained model
if not prepared and config.checkpoint_path is not '':
model = load_network(model, config.checkpoint_path, 'model')
# Return model
return model
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}")
def setup_dataset(config, mode, requirements, **kwargs):
"""
Create a dataset class
Parameters
----------
config : CfgNode
Configuration (cf. configs/default_config.py)
mode : str {'train', 'validation', 'test'}
Mode from which we want the dataset
requirements : dict (string -> bool)
Different requirements for dataset loading (gt_depth, gt_pose, etc)
kwargs : dict
Extra parameters for dataset creation
Returns
-------
dataset : Dataset
Dataset class for that mode
"""
# If no dataset is given, return None
if len(config.path) == 0:
return None
print0(pcolor('###### Setup %s datasets' % mode, 'red'))
# Global shared dataset arguments
dataset_args = {
'back_context': config.back_context,
'forward_context': config.forward_context,
'data_transform': get_transforms(mode, **kwargs),
'strides': config.strides
}
# Loop over all datasets
datasets = []
for i in range(len(config.split)):
path_split = os.path.join(config.path[i], config.split[i])
# Individual shared dataset arguments
dataset_args_i = {
'depth_type':
config.depth_type[i] if requirements['gt_depth'] else None,
'with_pose': requirements['gt_pose'],
}
# KITTI dataset
if config.dataset[i] == 'KITTI':
from dro_sfm.datasets.kitti_dataset import KITTIDataset
dataset = KITTIDataset(
config.path[i],
path_split,
**dataset_args,
**dataset_args_i,
)
# DGP dataset
elif config.dataset[i] == 'DGP':
from dro_sfm.datasets.dgp_dataset import DGPDataset
dataset = DGPDataset(
config.path[i],
config.split[i],
**dataset_args,
**dataset_args_i,
cameras=config.cameras[i],
)
# NYU dataset
elif config.dataset[i] == 'NYU':
from dro_sfm.datasets.nyu_dataset_processed import NYUDataset
dataset = NYUDataset(
config.path[i],
config.split[i],
**dataset_args,
**dataset_args_i,
)
# NYU dataset
elif config.dataset[i] == 'NYUtest':
from dro_sfm.datasets.nyu_dataset_test_processed import NYUDataset
dataset = NYUDataset(
config.path[i],
config.split[i],
**dataset_args,
**dataset_args_i,
)
# Demon dataset
elif config.dataset[i] == 'Demon':
from dro_sfm.datasets.demon_dataset import DemonDataset
dataset = DemonDataset(
config.path[i],
config.split[i],
**dataset_args,
**dataset_args_i,
)
# DemonMF dataset
elif config.dataset[i] == 'DemonMF':
from dro_sfm.datasets.demon_mf_dataset import DemonDataset
dataset = DemonDataset(
config.path[i],
config.split[i],
**dataset_args,
**dataset_args_i,
)
# Scannet dataset
elif config.dataset[i] == 'Scannet':
from dro_sfm.datasets.scannet_dataset import ScannetDataset
dataset = ScannetDataset(
config.path[i],
config.split[i],
**dataset_args,
**dataset_args_i,
)
# Scannet dataset
elif config.dataset[i] == 'ScannetTest':
from dro_sfm.datasets.scannet_test_dataset import ScannetTestDataset
dataset = ScannetTestDataset(
config.path[i],
config.split[i],
**dataset_args,
**dataset_args_i,
)
# Scannet dataset
elif config.dataset[i] == 'ScannetTestMF':
from dro_sfm.datasets.scannet_test_dataset_mf import ScannetTestDataset
dataset = ScannetTestDataset(
config.path[i],
config.split[i],
**dataset_args,
**dataset_args_i,
)
# Scannet banet dataset
elif config.dataset[i] == 'ScannetBA':
from dro_sfm.datasets.scannet_banet_dataset import ScannetBADataset
dataset = ScannetBADataset(
config.path[i],
config.split[i],
**dataset_args,
**dataset_args_i,
)
# Video dataset
elif config.dataset[i] == 'Video':
from dro_sfm.datasets.video_dataset import VideoDataset
dataset = VideoDataset(
config.path[i],
config.split[i],
**dataset_args,
**dataset_args_i,
)
# Video random sample dataset
elif config.dataset[i] == 'Video_Random':
from dro_sfm.datasets.video_random_dataset import VideoRandomDataset
dataset = VideoRandomDataset(
config.path[i],
config.split[i],
**dataset_args,
**dataset_args_i,
)
# Image dataset
elif config.dataset[i] == 'Image':
from dro_sfm.datasets.image_dataset import ImageDataset
dataset = ImageDataset(
config.path[i],
config.split[i],
**dataset_args,
**dataset_args_i,
)
else:
ValueError('Unknown dataset %d' % config.dataset[i])
# Repeat if needed
if 'repeat' in config and config.repeat[i] > 1:
dataset = ConcatDataset([dataset for _ in range(config.repeat[i])])
datasets.append(dataset)
# Display dataset information
bar = '######### {:>7}'.format(len(dataset))
if 'repeat' in config:
bar += ' (x{})'.format(config.repeat[i])
bar += ': {:<}'.format(path_split)
print0(pcolor(bar, 'yellow'))
# If training, concatenate all datasets into a single one
if mode == 'train':
datasets = [ConcatDataset(datasets)]
return datasets