def setup_depth_net(config, prepared, **kwargs): """ Create a depth 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("DepthNet: %s" % config.name, "yellow")) depth_net = load_class_args_create( config.name, paths=[ "packnet_sfm.networks.depth", ], args={ **config, **kwargs }, ) if not prepared and config.checkpoint_path is not "": depth_net = load_network(depth_net, config.checkpoint_path, ["depth_net", "disp_network"]) return depth_net
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')) model = load_class(config.name, paths=['packnet_sfm.models',])( **{**config.loss, **kwargs}) # Add depth network if required if 'depth_net' in model.network_requirements: model.add_depth_net(setup_depth_net(config.depth_net, prepared)) # Add pose network if required if 'pose_net' in model.network_requirements: model.add_pose_net(setup_pose_net(config.pose_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 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=[ 'packnet_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 main(args, N): # Initialize horovod hvd_init() # Parse arguments configs = [] state_dicts = [] for i in range(N): config, state_dict = parse_test_file(args.checkpoints[i]) configs.append(config) state_dicts.append(state_dict) # If no image shape is provided, use the checkpoint one image_shape = args.image_shape if image_shape is None: image_shape = configs[0].datasets.augmentation.image_shape # Set debug if requested set_debug(configs[0].debug) model_wrappers = [] for i in range(N): # Initialize model wrapper from checkpoint arguments model_wrappers.append(ModelWrapper(configs[i], load_datasets=False)) # Restore monodepth_model state model_wrappers[i].load_state_dict(state_dicts[i]) # 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(): for i in range(N): model_wrappers[i] = model_wrappers[i].to('cuda:{}'.format(rank()), dtype=dtype) # Set to eval mode for i in range(N): model_wrappers[i].eval() if args.input_folders is None: files = [[args.input_imgs[i]] for i in range(N)] else: files = [[] for i in range(N)] for i in range(N): for ext in ['png', 'jpg']: files[i] = glob.glob((os.path.join(args.input_folders[i], '*.{}'.format(ext)))) files[i].sort() print0('Found {} files'.format(len(files[i]))) n_files = len(files[0]) # Process each file infer_plot_and_save_3D_pcl(files, args.output, model_wrappers, image_shape, args.half, args.save, bool(int(args.stop)))
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 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, N): # Initialize horovod hvd_init() # Parse arguments configs = [] state_dicts = [] for i in range(N): config, state_dict = parse_test_file(args.checkpoints[i]) configs.append(config) state_dicts.append(state_dict) # If no image shape is provided, use the checkpoint one image_shape = args.image_shape if image_shape is None: image_shape = configs[0].datasets.augmentation.image_shape # Set debug if requested set_debug(configs[0].debug) model_wrappers = [] for i in range(N): # Initialize model wrapper from checkpoint arguments model_wrappers.append(ModelWrapper(configs[i], load_datasets=False)) # Restore monodepth_model state model_wrappers[i].load_state_dict(state_dicts[i]) # Send model to GPU if available if torch.cuda.is_available(): for i in range(N): model_wrappers[i] = model_wrappers[i].to('cuda:{}'.format(rank())) # Set to eval mode for i in range(N): model_wrappers[i].eval() if args.input_folder is None: files = [[args.input_imgs[i]] for i in range(N)] else: files = [[] for _ in range(N)] for i in range(N): for ext in ['png', 'jpg']: files[i] = glob( (os.path.join(args.input_folder, 'cam_' + str(i) + '/', '*.{}'.format(ext)))) files[i].sort() files[i] = files[i][::args.every_n_files] print0('Found {} files'.format(len(files[i]))) n_files = len(files[0]) # Process each file infer_optimal_calib(files, model_wrappers, image_shape)
def infer(ckpt_file, input_file, output_file, image_shape): """ Monocular depth estimation test script. Parameters ---------- ckpt_file : str Checkpoint path for a pretrained model input_file : str File or folder with input images output_file : str File or folder with output images image_shape : tuple Input image shape (H,W) """ # Initialize horovod hvd_init() # Parse arguments config, state_dict = parse_test_file(ckpt_file) # If no image shape is provided, use the checkpoint one 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) # Send model to GPU if available if torch.cuda.is_available(): model_wrapper = model_wrapper.to('cuda:{}'.format(rank())) if os.path.isdir(input_file): # If input file is a folder, search for image files files = [] for ext in ['png', 'jpg']: files.extend(glob((os.path.join(input_file, '*.{}'.format(ext))))) files.sort() print0('Found {} files'.format(len(files))) else: # Otherwise, use it as is files = [input_file] # Process each file for file in files[rank()::world_size()]: process(file, output_file, model_wrapper, image_shape)
def prepare_datasets(self): """Prepare datasets for training, validation and test.""" # Prepare datasets print0(pcolor('### Preparing Datasets', 'green')) augmentation = self.config.datasets.augmentation self.train_dataset = setup_dataset(self.config.datasets.train, 'train', self.model.requires_gt_depth, **augmentation) self.validation_dataset = setup_dataset( self.config.datasets.validation, 'validation', **augmentation) self.test_dataset = setup_dataset(self.config.datasets.test, 'test', **augmentation)
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 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 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_plot_and_save_3D_pcl(fn, args.output, model_wrapper, image_shape, args.half, args.save)
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')) # SfmModel, SelfSupModel, VelSupModel loaded model = load_class(config.name, paths=['packnet_sfm.models',])( **{**config.loss, **kwargs}) # Add depth network if required if model.network_requirements['depth_net']: model.add_depth_net(setup_depth_net(config.depth_net, prepared, num_scales=config.loss.num_scales, min_depth=config.params.min_depth, max_depth=config.params.max_depth, upsample_depth_maps=config.loss.upsample_depth_maps )) # Add pose network if required if model.network_requirements['pose_net']: model.add_pose_net( setup_pose_net(config.pose_net, prepared, rotation_mode=config.loss.rotation_mode, **kwargs)) # 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 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_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, 'with_geometric_context': config.with_geometric_context, } # 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 if config.dataset[i] == 'ValeoMultifocal': dataset_args_i = { 'depth_type': config.depth_type[i] if requirements['gt_depth'] else None, 'with_pose': requirements['gt_pose'], 'data_transform': get_transforms_multifocal(mode, **kwargs), 'with_spatiotemp_context': config.with_spatiotemp_context, } elif config.dataset[i] == 'KITTIValeoFisheye': dataset_args_i = { 'depth_type': config.depth_type[i] if requirements['gt_depth'] else None, 'with_pose': requirements['gt_pose'], 'data_transform': get_transforms_fisheye(mode, **kwargs), 'calibrations_suffix': config.calibrations_suffix, 'depth_suffix': config.depth_suffix, 'cam_convs': config.cam_convs } elif config.dataset[i] == 'KITTIValeoDistorted': dataset_args_i = { 'depth_type': config.depth_type[i] if requirements['gt_depth'] else None, 'with_pose': requirements['gt_pose'], 'data_transform': get_transforms_distorted(mode, **kwargs) } elif config.dataset[i] == 'DGPvaleo': dataset_args_i = { 'depth_type': config.depth_type[i] if requirements['gt_depth'] else None, 'with_pose': requirements['gt_pose'], 'data_transform': get_transforms_dgp_valeo(mode, **kwargs) } elif config.dataset[i] == 'WoodscapeFisheye': dataset_args_i = { 'depth_type': config.depth_type[i] if requirements['gt_depth'] else None, 'with_pose': requirements['gt_pose'], 'data_transform': get_transforms_woodscape_fisheye(mode, **kwargs) } else: dataset_args_i = { 'depth_type': config.depth_type[i] if requirements['gt_depth'] else None, 'with_pose': requirements['gt_pose'], 'data_transform': get_transforms(mode, **kwargs) } if config.dataset[i] == 'ValeoMultifocal': from packnet_sfm.datasets.kitti_based_valeo_dataset_multifocal import KITTIBasedValeoDatasetMultifocal dataset = KITTIBasedValeoDatasetMultifocal( config.path[i], path_split, **dataset_args, **dataset_args_i, cameras=config.cameras[i], ) # KITTI dataset elif config.dataset[i] == 'KITTI': from packnet_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 packnet_sfm.datasets.dgp_dataset import DGPDataset dataset = DGPDataset( config.path[i], config.split[i], **dataset_args, **dataset_args_i, cameras=config.cameras[i], ) # DGP dataset elif config.dataset[i] == 'DGPvaleo': from packnet_sfm.datasets.dgp_valeo_dataset import DGPvaleoDataset dataset = DGPvaleoDataset( config.path[i], config.split[i], **dataset_args, **dataset_args_i, cameras=config.cameras[i], ) # Image dataset elif config.dataset[i] == 'Image': from packnet_sfm.datasets.image_dataset import ImageDataset dataset = ImageDataset( config.path[i], config.split[i], **dataset_args, **dataset_args_i, ) # KITTI-based Valeo dataset elif config.dataset[i] == 'KITTIValeo': from packnet_sfm.datasets.kitti_based_valeo_dataset import KITTIBasedValeoDataset dataset = KITTIBasedValeoDataset( config.path[i], path_split, **dataset_args, **dataset_args_i, cameras=config.cameras[i], ) # KITTI-based Valeo dataset (fisheye) elif config.dataset[i] == 'KITTIValeoFisheye': from packnet_sfm.datasets.kitti_based_valeo_dataset_fisheye_singleView import \ KITTIBasedValeoDatasetFisheye_singleView dataset = KITTIBasedValeoDatasetFisheye_singleView( config.path[i], path_split, **dataset_args, **dataset_args_i, cameras=config.cameras[i], ) elif config.dataset[i] == 'KITTIValeoDistorted': from packnet_sfm.datasets.kitti_based_valeo_dataset_distorted_singleView import \ KITTIBasedValeoDatasetDistorted_singleView dataset = KITTIBasedValeoDatasetDistorted_singleView( config.path[i], path_split, **dataset_args, **dataset_args_i, cameras=config.cameras[i], ) elif config.dataset[i] == 'WoodscapeFisheye': from packnet_sfm.datasets.woodscape_fisheye import WoodscapeFisheye dataset = WoodscapeFisheye( config.path[i], path_split, **dataset_args, **dataset_args_i, cameras=config.cameras[i], ) # Image-based Valeo dataset elif config.dataset[i] == 'ImageValeo': from packnet_sfm.datasets.image_based_valeo_dataset import ImageBasedValeoDataset dataset = ImageBasedValeoDataset( config.path[i], config.split[i], **dataset_args, **dataset_args_i, cameras=config.cameras[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
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), } # 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 packnet_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 packnet_sfm.datasets.dgp_dataset import DGPDataset dataset = DGPDataset( config.path[i], config.split[i], **dataset_args, **dataset_args_i, cameras=config.cameras[i], ) # Image dataset elif config.dataset[i] == "Image": from packnet_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
def main(args): # Initialize horovod hvd_init() # Parse arguments config1, state_dict1 = parse_test_file(args.checkpoint1) config2, state_dict2 = parse_test_file(args.checkpoint2) config3, state_dict3 = parse_test_file(args.checkpoint3) config4, state_dict4 = parse_test_file(args.checkpoint4) # If no image shape is provided, use the checkpoint one image_shape = args.image_shape if image_shape is None: image_shape = config1.datasets.augmentation.image_shape # Set debug if requested set_debug(config1.debug) # Initialize model wrapper from checkpoint arguments model_wrapper1 = ModelWrapper(config1, load_datasets=False) model_wrapper2 = ModelWrapper(config2, load_datasets=False) model_wrapper3 = ModelWrapper(config3, load_datasets=False) model_wrapper4 = ModelWrapper(config4, load_datasets=False) # Restore monodepth_model state model_wrapper1.load_state_dict(state_dict1) model_wrapper2.load_state_dict(state_dict2) model_wrapper3.load_state_dict(state_dict3) model_wrapper4.load_state_dict(state_dict4) # 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_wrapper1 = model_wrapper1.to('cuda:{}'.format(rank()), dtype=dtype) model_wrapper2 = model_wrapper2.to('cuda:{}'.format(rank()), dtype=dtype) model_wrapper3 = model_wrapper3.to('cuda:{}'.format(rank()), dtype=dtype) model_wrapper4 = model_wrapper4.to('cuda:{}'.format(rank()), dtype=dtype) # Set to eval mode model_wrapper1.eval() model_wrapper2.eval() model_wrapper3.eval() model_wrapper4.eval() if os.path.isdir(args.input1): # If input file is a folder, search for image files files1 = [] for ext in ['png', 'jpg']: files1.extend(glob((os.path.join(args.input1, '*.{}'.format(ext))))) files1.sort() print0('Found {} files'.format(len(files1))) else: # Otherwise, use it as is files1 = [args.input1] if os.path.isdir(args.input2): # If input file is a folder, search for image files files2 = [] for ext in ['png', 'jpg']: files2.extend(glob((os.path.join(args.input2, '*.{}'.format(ext))))) files2.sort() print0('Found {} files'.format(len(files2))) else: # Otherwise, use it as is files2 = [args.input2] if os.path.isdir(args.input3): # If input file is a folder, search for image files files3 = [] for ext in ['png', 'jpg']: files3.extend(glob((os.path.join(args.input3, '*.{}'.format(ext))))) files3.sort() print0('Found {} files'.format(len(files3))) else: # Otherwise, use it as is files3 = [args.input3] if os.path.isdir(args.input4): # If input file is a folder, search for image files files4 = [] for ext in ['png', 'jpg']: files4.extend(glob((os.path.join(args.input4, '*.{}'.format(ext))))) files4.sort() print0('Found {} files'.format(len(files4))) else: # Otherwise, use it as is files4 = [args.input4] n_files = len(files1) # Process each file for fn1, fn2, fn3, fn4 in zip(files1[rank()::world_size()], files2[rank()::world_size()], files3[rank()::world_size()], files4[rank()::world_size()]): infer_plot_and_save_3D_pcl(fn1, fn2, fn3, fn4, args.output1, args.output2, args.output3, args.output4, model_wrapper1, model_wrapper2, model_wrapper3, model_wrapper4, bool(int(args.hasGTdepth1)), bool(int(args.hasGTdepth2)), bool(int(args.hasGTdepth3)), bool(int(args.hasGTdepth4)), image_shape, args.half, args.save)
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) } # 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 'gt_depth' in requirements else None, 'input_depth_type': config.input_depth_type[i] if 'gt_depth' in requirements else None, 'with_pose': 'gt_pose' in requirements, } # KITTI dataset if config.dataset[i] == 'KITTI': from packnet_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 packnet_sfm.datasets.dgp_dataset import DGPDataset dataset = DGPDataset( config.path[i], config.split[i], **dataset_args, **dataset_args_i, cameras=config.cameras[i], ) # Image dataset elif config.dataset[i] == 'Image': from packnet_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
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: 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) f = open(args.output + ".txt", 'w') for key in sorted(poses.keys()): rot_matrix, translation = poses[key] # print(rot_matrix, translation) # print("orientation, position") orientation = orientation.dot(rot_matrix.tolist()) position += orientation.dot(translation.tolist()) # print(torch.tensor(orientation)) q = transforms.matrix_to_quaternion(torch.tensor(orientation)) q = q.numpy() # print(q[0]) # print(position) f.write("%.10f %.10f %.10f %.10f %.10f %.10f %.10f\n" % (position[0], position[1], position[2], q[0][3], q[0][2], q[0][1], q[0][0])) # f.write("{.10f} {.10f} {.10f} {.10f} {.10f} {.10f} {.10f}" # .format(position[0], position[1], position[2], q[0][1], q[0][2], q[0][3], q[0][0])) # 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]} f.close() # json.dump(poses, open(args.output, "w"), sort_keys=True) print(f"Written pose of {len(list_of_files)} images to {args.output}")