def _sync_s3(self, filepath, model):
# If it's not time to sync, do nothing
if self.s3_enabled and (model.current_epoch +
1) % self.s3_frequency == 0:
filepath = os.path.dirname(filepath)
# Print message and links
print(
pcolor(
"###### Syncing: {} -> {}".format(
filepath, model.config.checkpoint.s3_path),
"red",
attrs=["bold"],
))
print(
pcolor(
"###### URL: {}".format(model.config.checkpoint.s3_url),
"red",
attrs=["bold"],
))
# If it's time to save code
if self.save_code:
self.save_code = False
save_code(filepath)
# Sync model to s3
sync_s3_data(filepath, model)
def infer_and_save_depth(input_file, output_file, model_wrapper, image_shape,
half, save):
"""
Process a single input file to produce and save visualization
Parameters
----------
input_file : str
Image file
output_file : str
Output file, or folder where the output will be saved
model_wrapper : nn.Module
Model wrapper used for inference
image_shape : Image shape
Input image shape
half: bool
use half precision (fp16)
save: str
Save format (npz or png)
"""
if not is_image(output_file):
# If not an image, assume it's a folder and append the input name
os.makedirs(output_file, exist_ok=True)
output_file = os.path.join(output_file, os.path.basename(input_file))
# change to half precision for evaluation if requested
dtype = torch.float16 if half else None
# Load image
image = load_image(input_file)
# Resize and to tensor
image = resize_image(image, image_shape)
image = to_tensor(image).unsqueeze(0)
# Send image to GPU if available
if torch.cuda.is_available():
image = image.to('cuda:{}'.format(rank()), dtype=dtype)
# Depth inference (returns predicted inverse depth)
pred_inv_depth = model_wrapper.depth(image)[0]
if save == 'npz' or save == 'png':
# Get depth from predicted depth map and save to different formats
filename = '{}.{}'.format(os.path.splitext(output_file)[0], save)
print('Saving {} to {}'.format(
pcolor(input_file, 'cyan', attrs=['bold']),
pcolor(filename, 'magenta', attrs=['bold'])))
write_depth(filename, depth=inv2depth(pred_inv_depth))
else:
# Prepare RGB image
rgb = image[0].permute(1, 2, 0).detach().cpu().numpy() * 255
# Prepare inverse depth
viz_pred_inv_depth = viz_inv_depth(pred_inv_depth[0]) * 255
# Concatenate both vertically
image = np.concatenate([rgb, viz_pred_inv_depth], 0)
# Save visualization
print('Saving {} to {}'.format(
pcolor(input_file, 'cyan', attrs=['bold']),
pcolor(output_file, 'magenta', attrs=['bold'])))
imwrite(output_file, image[:, :, ::-1])
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 print_metrics(self, metrics_data, dataset):
"""Print depth metrics on rank 0 if available"""
if not metrics_data[0]:
return
hor_line = '|{:<}|'.format('*' * 93)
met_line = '| {:^14} | {:^8} | {:^8} | {:^8} | {:^8} | {:^8} | {:^8} | {:^8} |'
num_line = '{:<14} | {:^8.3f} | {:^8.3f} | {:^8.3f} | {:^8.3f} | {:^8.3f} | {:^8.3f} | {:^8.3f}'
def wrap(string):
return '| {} |'.format(string)
print()
print()
print()
print(hor_line)
if self.optimizer is not None:
bs = 'E: {} BS: {}'.format(self.current_epoch + 1,
self.config.datasets.train.batch_size)
if self.model is not None:
bs += ' - {}'.format(self.config.model.name)
lr = 'LR ({}):'.format(self.config.model.optimizer.name)
for param in self.optimizer.param_groups:
lr += ' {} {:.2e}'.format(param['name'], param['lr'])
par_line = wrap(pcolor('{:<40}{:>51}'.format(bs, lr),
'green', attrs=['bold', 'dark']))
print(par_line)
print(hor_line)
print(met_line.format(*(('METRIC',) + self.metrics_keys)))
for n, metrics in enumerate(metrics_data):
print(hor_line)
path_line = '{}'.format(
os.path.join(dataset.path[n], dataset.split[n]))
if len(dataset.cameras[n]) == 1: # only allows single cameras
path_line += ' ({})'.format(dataset.cameras[n][0])
print(wrap(pcolor('*** {:<87}'.format(path_line), 'magenta', attrs=['bold'])))
print(hor_line)
for key, metric in metrics.items():
if self.metrics_name in key:
print(wrap(pcolor(num_line.format(
*((key.upper(),) + tuple(metric.tolist()))), 'cyan')))
print(hor_line)
if self.logger:
run_line = wrap(pcolor('{:<60}{:>31}'.format(
self.config.wandb.url, self.config.wandb.name), 'yellow', attrs=['dark']))
print(run_line)
print(hor_line)
print()
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 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 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 process(input_file, output_file, model_wrapper, image_shape):
"""
Process a single input file to produce and save visualization
Parameters
----------
input_file : str
Image file
output_file : str
Output file, or folder where the output will be saved
model_wrapper : nn.Module
Model wrapper used for inference
image_shape : Image shape
Input image shape
Returns
-------
"""
# Load image
image = load_image(input_file)
# Resize and to tensor
image = resize_image(image, image_shape)
image = to_tensor(image).unsqueeze(0)
# Send image to GPU if available
if torch.cuda.is_available():
image = image.to('cuda:{}'.format(rank()))
# Depth inference
depth = model_wrapper.depth(image)[0]
# Prepare RGB image
rgb_i = image[0].permute(1, 2, 0).detach().cpu().numpy() * 255
# Prepare inverse depth
pred_inv_depth_i = viz_inv_depth(depth[0]) * 255
# Concatenate both vertically
image = np.concatenate([rgb_i, pred_inv_depth_i], 0)
if not is_image(output_file):
# If not an image, assume it's a folder and append the input name
os.makedirs(output_file, exist_ok=True)
output_file = os.path.join(output_file, os.path.basename(input_file))
# Save visualization
print('Saving {} to {}'.format(
pcolor(input_file, 'cyan', attrs=['bold']),
pcolor(output_file, 'magenta', attrs=['bold'])))
imwrite(output_file, image[:, :, ::-1])
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_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'))
# 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 print_metrics(self, metrics_data, dataset):
"""Print depth metrics on rank 0 if available"""
if not metrics_data[0]:
return
hor_line = "|{:<}|".format("*" * 93)
met_line = "| {:^14} | {:^8} | {:^8} | {:^8} | {:^8} | {:^8} | {:^8} | {:^8} |"
num_line = "{:<14} | {:^8.3f} | {:^8.3f} | {:^8.3f} | {:^8.3f} | {:^8.3f} | {:^8.3f} | {:^8.3f}"
def wrap(string):
return "| {} |".format(string)
print()
print()
print()
print(hor_line)
if self.optimizer is not None:
bs = "E: {} BS: {}".format(self.current_epoch + 1,
self.config.datasets.train.batch_size)
if self.model is not None:
bs += " - {}".format(self.config.model.name)
lr = "LR ({}):".format(self.config.model.optimizer.name)
for param in self.optimizer.param_groups:
lr += " {} {:.2e}".format(param["name"], param["lr"])
par_line = wrap(
pcolor("{:<40}{:>51}".format(bs, lr),
"green",
attrs=["bold", "dark"]))
print(par_line)
print(hor_line)
print(met_line.format(*(("METRIC", ) + self.metrics_keys)))
for n, metrics in enumerate(metrics_data):
print(hor_line)
path_line = "{}".format(
os.path.join(dataset.path[n], dataset.split[n]))
if len(dataset.cameras[n]) == 1: # only allows single cameras
path_line += " ({})".format(dataset.cameras[n][0])
print(
wrap(
pcolor("*** {:<87}".format(path_line),
"magenta",
attrs=["bold"])))
print(hor_line)
for key, metric in metrics.items():
if self.metrics_name in key:
print(
wrap(
pcolor(
num_line.format(*((key.upper(), ) +
tuple(metric.tolist()))),
"cyan",
)))
print(hor_line)
if self.logger:
run_line = wrap(
pcolor(
"{:<60}{:>31}".format(self.config.wandb.url,
self.config.wandb.name),
"yellow",
attrs=["dark"],
))
print(run_line)
print(hor_line)
print()
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 infer_plot_and_save_3D_pcl(input_file1, input_file2, input_file3,
input_file4, output_file1, output_file2,
output_file3, output_file4, model_wrapper1,
model_wrapper2, model_wrapper3, model_wrapper4,
hasGTdepth1, hasGTdepth2, hasGTdepth3,
hasGTdepth4, image_shape, half, save):
"""
Process a single input file to produce and save visualization
Parameters
----------
input_file : str
Image file
output_file : str
Output file, or folder where the output will be saved
model_wrapper : nn.Module
Model wrapper used for inference
image_shape : Image shape
Input image shape
half: bool
use half precision (fp16)
save: str
Save format (npz or png)
"""
if not is_image(output_file1):
# If not an image, assume it's a folder and append the input name
os.makedirs(output_file1, exist_ok=True)
output_file1 = os.path.join(output_file1,
os.path.basename(input_file1))
if not is_image(output_file2):
# If not an image, assume it's a folder and append the input name
os.makedirs(output_file2, exist_ok=True)
output_file2 = os.path.join(output_file2,
os.path.basename(input_file2))
if not is_image(output_file3):
# If not an image, assume it's a folder and append the input name
os.makedirs(output_file3, exist_ok=True)
output_file3 = os.path.join(output_file3,
os.path.basename(input_file3))
if not is_image(output_file4):
# If not an image, assume it's a folder and append the input name
os.makedirs(output_file4, exist_ok=True)
output_file4 = os.path.join(output_file4,
os.path.basename(input_file4))
# change to half precision for evaluation if requested
dtype = torch.float16 if half else None
# Load image
image1 = load_image(input_file1).convert('RGB')
image2 = load_image(input_file2).convert('RGB')
image3 = load_image(input_file3).convert('RGB')
image4 = load_image(input_file4).convert('RGB')
# Resize and to tensor
image1 = resize_image(image1, image_shape)
image2 = resize_image(image2, image_shape)
image3 = resize_image(image3, image_shape)
image4 = resize_image(image4, image_shape)
image1 = to_tensor(image1).unsqueeze(0)
image2 = to_tensor(image2).unsqueeze(0)
image3 = to_tensor(image3).unsqueeze(0)
image4 = to_tensor(image4).unsqueeze(0)
# Send image to GPU if available
if torch.cuda.is_available():
image1 = image1.to('cuda:{}'.format(rank()), dtype=dtype)
image2 = image2.to('cuda:{}'.format(rank()), dtype=dtype)
image3 = image3.to('cuda:{}'.format(rank()), dtype=dtype)
image4 = image4.to('cuda:{}'.format(rank()), dtype=dtype)
# Depth inference (returns predicted inverse depth)
pred_inv_depth1 = model_wrapper1.depth(image1)
pred_inv_depth2 = model_wrapper2.depth(image2)
pred_inv_depth3 = model_wrapper1.depth(image3)
pred_inv_depth4 = model_wrapper2.depth(image4)
pred_depth1 = inv2depth(pred_inv_depth1)
pred_depth2 = inv2depth(pred_inv_depth2)
pred_depth3 = inv2depth(pred_inv_depth3)
pred_depth4 = inv2depth(pred_inv_depth4)
base_folder_str1 = get_base_folder(input_file1)
split_type_str1 = get_split_type(input_file1)
seq_name_str1 = get_sequence_name(input_file1)
camera_str1 = get_camera_name(input_file1)
base_folder_str2 = get_base_folder(input_file2)
split_type_str2 = get_split_type(input_file2)
seq_name_str2 = get_sequence_name(input_file2)
camera_str2 = get_camera_name(input_file2)
base_folder_str3 = get_base_folder(input_file3)
split_type_str3 = get_split_type(input_file3)
seq_name_str3 = get_sequence_name(input_file3)
camera_str3 = get_camera_name(input_file3)
base_folder_str4 = get_base_folder(input_file4)
split_type_str4 = get_split_type(input_file4)
seq_name_str4 = get_sequence_name(input_file4)
camera_str4 = get_camera_name(input_file4)
calib_data1 = {}
calib_data2 = {}
calib_data3 = {}
calib_data4 = {}
calib_data1[camera_str1] = read_raw_calib_files_camera_valeo(
base_folder_str1, split_type_str1, seq_name_str1, camera_str1)
calib_data2[camera_str2] = read_raw_calib_files_camera_valeo(
base_folder_str2, split_type_str2, seq_name_str2, camera_str2)
calib_data3[camera_str3] = read_raw_calib_files_camera_valeo(
base_folder_str3, split_type_str3, seq_name_str3, camera_str3)
calib_data4[camera_str4] = read_raw_calib_files_camera_valeo(
base_folder_str4, split_type_str4, seq_name_str4, camera_str4)
path_to_theta_lut1 = get_path_to_theta_lut(input_file1)
path_to_ego_mask1 = get_path_to_ego_mask(input_file1)
poly_coeffs1, principal_point1, scale_factors1 = get_intrinsics(
input_file1, calib_data1)
path_to_theta_lut2 = get_path_to_theta_lut(input_file2)
path_to_ego_mask2 = get_path_to_ego_mask(input_file2)
poly_coeffs2, principal_point2, scale_factors2 = get_intrinsics(
input_file2, calib_data2)
path_to_theta_lut3 = get_path_to_theta_lut(input_file3)
path_to_ego_mask3 = get_path_to_ego_mask(input_file3)
poly_coeffs3, principal_point3, scale_factors3 = get_intrinsics(
input_file3, calib_data3)
path_to_theta_lut4 = get_path_to_theta_lut(input_file4)
path_to_ego_mask4 = get_path_to_ego_mask(input_file4)
poly_coeffs4, principal_point4, scale_factors4 = get_intrinsics(
input_file4, calib_data4)
poly_coeffs1 = torch.from_numpy(poly_coeffs1).unsqueeze(0)
principal_point1 = torch.from_numpy(principal_point1).unsqueeze(0)
scale_factors1 = torch.from_numpy(scale_factors1).unsqueeze(0)
poly_coeffs2 = torch.from_numpy(poly_coeffs2).unsqueeze(0)
principal_point2 = torch.from_numpy(principal_point2).unsqueeze(0)
scale_factors2 = torch.from_numpy(scale_factors2).unsqueeze(0)
poly_coeffs3 = torch.from_numpy(poly_coeffs3).unsqueeze(0)
principal_point3 = torch.from_numpy(principal_point3).unsqueeze(0)
scale_factors3 = torch.from_numpy(scale_factors3).unsqueeze(0)
poly_coeffs4 = torch.from_numpy(poly_coeffs4).unsqueeze(0)
principal_point4 = torch.from_numpy(principal_point4).unsqueeze(0)
scale_factors4 = torch.from_numpy(scale_factors4).unsqueeze(0)
pose_matrix1 = torch.from_numpy(
get_extrinsics_pose_matrix(input_file1, calib_data1)).unsqueeze(0)
pose_matrix2 = torch.from_numpy(
get_extrinsics_pose_matrix(input_file2, calib_data2)).unsqueeze(0)
pose_matrix3 = torch.from_numpy(
get_extrinsics_pose_matrix(input_file3, calib_data3)).unsqueeze(0)
pose_matrix4 = torch.from_numpy(
get_extrinsics_pose_matrix(input_file4, calib_data4)).unsqueeze(0)
pose_tensor1 = Pose(pose_matrix1)
pose_tensor2 = Pose(pose_matrix2)
pose_tensor3 = Pose(pose_matrix3)
pose_tensor4 = Pose(pose_matrix4)
ego_mask1 = np.load(path_to_ego_mask1)
ego_mask2 = np.load(path_to_ego_mask2)
ego_mask3 = np.load(path_to_ego_mask3)
ego_mask4 = np.load(path_to_ego_mask4)
not_masked1 = ego_mask1.astype(bool).reshape(-1)
not_masked2 = ego_mask2.astype(bool).reshape(-1)
not_masked3 = ego_mask3.astype(bool).reshape(-1)
not_masked4 = ego_mask4.astype(bool).reshape(-1)
cam1 = CameraFisheye(path_to_theta_lut=[path_to_theta_lut1],
path_to_ego_mask=[path_to_ego_mask1],
poly_coeffs=poly_coeffs1.float(),
principal_point=principal_point1.float(),
scale_factors=scale_factors1.float(),
Tcw=pose_tensor1)
cam2 = CameraFisheye(path_to_theta_lut=[path_to_theta_lut2],
path_to_ego_mask=[path_to_ego_mask2],
poly_coeffs=poly_coeffs2.float(),
principal_point=principal_point2.float(),
scale_factors=scale_factors2.float(),
Tcw=pose_tensor2)
cam3 = CameraFisheye(path_to_theta_lut=[path_to_theta_lut3],
path_to_ego_mask=[path_to_ego_mask3],
poly_coeffs=poly_coeffs3.float(),
principal_point=principal_point3.float(),
scale_factors=scale_factors3.float(),
Tcw=pose_tensor3)
cam4 = CameraFisheye(path_to_theta_lut=[path_to_theta_lut4],
path_to_ego_mask=[path_to_ego_mask4],
poly_coeffs=poly_coeffs4.float(),
principal_point=principal_point4.float(),
scale_factors=scale_factors4.float(),
Tcw=pose_tensor4)
if torch.cuda.is_available():
cam1 = cam1.to('cuda:{}'.format(rank()), dtype=dtype)
cam2 = cam2.to('cuda:{}'.format(rank()), dtype=dtype)
cam3 = cam3.to('cuda:{}'.format(rank()), dtype=dtype)
cam4 = cam4.to('cuda:{}'.format(rank()), dtype=dtype)
world_points1 = cam1.reconstruct(pred_depth1, frame='w')
world_points1 = world_points1[0].cpu().numpy()
world_points1 = world_points1.reshape((3, -1)).transpose()
world_points2 = cam2.reconstruct(pred_depth2, frame='w')
world_points2 = world_points2[0].cpu().numpy()
world_points2 = world_points2.reshape((3, -1)).transpose()
world_points3 = cam3.reconstruct(pred_depth3, frame='w')
world_points3 = world_points3[0].cpu().numpy()
world_points3 = world_points3.reshape((3, -1)).transpose()
world_points4 = cam4.reconstruct(pred_depth4, frame='w')
world_points4 = world_points4[0].cpu().numpy()
world_points4 = world_points4.reshape((3, -1)).transpose()
if hasGTdepth1:
gt_depth_file1 = get_depth_file(input_file1)
gt_depth1 = np.load(gt_depth_file1)['velodyne_depth'].astype(
np.float32)
gt_depth1 = torch.from_numpy(gt_depth1).unsqueeze(0).unsqueeze(0)
if torch.cuda.is_available():
gt_depth1 = gt_depth1.to('cuda:{}'.format(rank()), dtype=dtype)
gt_depth_3d1 = cam1.reconstruct(gt_depth1, frame='w')
gt_depth_3d1 = gt_depth_3d1[0].cpu().numpy()
gt_depth_3d1 = gt_depth_3d1.reshape((3, -1)).transpose()
if hasGTdepth2:
gt_depth_file2 = get_depth_file(input_file2)
gt_depth2 = np.load(gt_depth_file2)['velodyne_depth'].astype(
np.float32)
gt_depth2 = torch.from_numpy(gt_depth2).unsqueeze(0).unsqueeze(0)
if torch.cuda.is_available():
gt_depth2 = gt_depth2.to('cuda:{}'.format(rank()), dtype=dtype)
gt_depth_3d2 = cam2.reconstruct(gt_depth2, frame='w')
gt_depth_3d2 = gt_depth_3d2[0].cpu().numpy()
gt_depth_3d2 = gt_depth_3d2.reshape((3, -1)).transpose()
if hasGTdepth3:
gt_depth_file3 = get_depth_file(input_file3)
gt_depth3 = np.load(gt_depth_file3)['velodyne_depth'].astype(
np.float33)
gt_depth3 = torch.from_numpy(gt_depth3).unsqueeze(0).unsqueeze(0)
if torch.cuda.is_available():
gt_depth3 = gt_depth3.to('cuda:{}'.format(rank()), dtype=dtype)
gt_depth_3d3 = cam3.reconstruct(gt_depth3, frame='w')
gt_depth_3d3 = gt_depth_3d3[0].cpu().numpy()
gt_depth_3d3 = gt_depth_3d3.reshape((3, -1)).transpose()
if hasGTdepth4:
gt_depth_file4 = get_depth_file(input_file4)
gt_depth4 = np.load(gt_depth_file4)['velodyne_depth'].astype(
np.float34)
gt_depth4 = torch.from_numpy(gt_depth4).unsqueeze(0).unsqueeze(0)
if torch.cuda.is_available():
gt_depth4 = gt_depth4.to('cuda:{}'.format(rank()), dtype=dtype)
gt_depth_3d4 = cam4.reconstruct(gt_depth4, frame='w')
gt_depth_3d4 = gt_depth_3d4[0].cpu().numpy()
gt_depth_3d4 = gt_depth_3d4.reshape((3, -1)).transpose()
world_points1 = world_points1[not_masked1]
world_points2 = world_points2[not_masked2]
world_points3 = world_points3[not_masked3]
world_points4 = world_points4[not_masked4]
if hasGTdepth1:
gt_depth_3d1 = gt_depth_3d1[not_masked1]
if hasGTdepth2:
gt_depth_3d2 = gt_depth_3d2[not_masked1]
if hasGTdepth3:
gt_depth_3d3 = gt_depth_3d3[not_masked3]
if hasGTdepth4:
gt_depth_3d4 = gt_depth_3d4[not_masked3]
pcl1 = o3d.geometry.PointCloud()
pcl1.points = o3d.utility.Vector3dVector(world_points1)
img_numpy1 = image1[0].cpu().numpy()
img_numpy1 = img_numpy1.reshape((3, -1)).transpose()
pcl2 = o3d.geometry.PointCloud()
pcl2.points = o3d.utility.Vector3dVector(world_points2)
img_numpy2 = image2[0].cpu().numpy()
img_numpy2 = img_numpy2.reshape((3, -1)).transpose()
pcl3 = o3d.geometry.PointCloud()
pcl3.points = o3d.utility.Vector3dVector(world_points3)
img_numpy3 = image3[0].cpu().numpy()
img_numpy3 = img_numpy3.reshape((3, -1)).transpose()
pcl4 = o3d.geometry.PointCloud()
pcl4.points = o3d.utility.Vector3dVector(world_points4)
img_numpy4 = image4[0].cpu().numpy()
img_numpy4 = img_numpy4.reshape((3, -1)).transpose()
img_numpy1 = img_numpy1[not_masked1]
pcl1.colors = o3d.utility.Vector3dVector(img_numpy1)
img_numpy2 = img_numpy2[not_masked2]
pcl2.colors = o3d.utility.Vector3dVector(img_numpy2)
img_numpy3 = img_numpy3[not_masked3]
pcl3.colors = o3d.utility.Vector3dVector(img_numpy3)
img_numpy4 = img_numpy4[not_masked4]
pcl4.colors = o3d.utility.Vector3dVector(img_numpy4)
#pcl.paint_uniform_color([1.0, 0.0, 0])
#print("Radius oulier removal")
#cl, ind = pcl.remove_radius_outlier(nb_points=10, radius=0.5)
#display_inlier_outlier(pcl, ind)
remove_outliers = True
if remove_outliers:
cl1, ind1 = pcl1.remove_statistical_outlier(nb_neighbors=10,
std_ratio=1.3)
inlier_cloud1 = pcl1.select_by_index(ind1)
outlier_cloud1 = pcl1.select_by_index(ind1, invert=True)
outlier_cloud1.paint_uniform_color([0.0, 0.0, 1.0])
cl2, ind2 = pcl2.remove_statistical_outlier(nb_neighbors=10,
std_ratio=1.3)
inlier_cloud2 = pcl2.select_by_index(ind2)
outlier_cloud2 = pcl2.select_by_index(ind2, invert=True)
outlier_cloud2.paint_uniform_color([0.0, 0.0, 1.0])
cl3, ind3 = pcl3.remove_statistical_outlier(nb_neighbors=10,
std_ratio=1.3)
inlier_cloud3 = pcl3.select_by_index(ind3)
outlier_cloud3 = pcl3.select_by_index(ind3, invert=True)
outlier_cloud3.paint_uniform_color([0.0, 0.0, 1.0])
cl4, ind4 = pcl4.remove_statistical_outlier(nb_neighbors=10,
std_ratio=1.3)
inlier_cloud4 = pcl4.select_by_index(ind4)
outlier_cloud4 = pcl4.select_by_index(ind4, invert=True)
outlier_cloud4.paint_uniform_color([0.0, 0.0, 1.0])
if hasGTdepth1:
pcl_gt1 = o3d.geometry.PointCloud()
pcl_gt1.points = o3d.utility.Vector3dVector(gt_depth_3d1)
pcl_gt1.paint_uniform_color([1.0, 0.0, 0])
if hasGTdepth2:
pcl_gt2 = o3d.geometry.PointCloud()
pcl_gt2.points = o3d.utility.Vector3dVector(gt_depth_3d2)
pcl_gt2.paint_uniform_color([1.0, 0.0, 0])
if hasGTdepth3:
pcl_gt3 = o3d.geometry.PointCloud()
pcl_gt3.points = o3d.utility.Vector3dVector(gt_depth_3d3)
pcl_gt3.paint_uniform_color([1.0, 0.0, 0])
if hasGTdepth4:
pcl_gt4 = o3d.geometry.PointCloud()
pcl_gt4.points = o3d.utility.Vector3dVector(gt_depth_3d4)
pcl_gt4.paint_uniform_color([1.0, 0.0, 0])
if remove_outliers:
toPlot = [inlier_cloud1, inlier_cloud2, inlier_cloud3, inlier_cloud4]
if hasGTdepth1:
toPlot.append(pcl_gt1)
if hasGTdepth2:
toPlot.append(pcl_gt2)
if hasGTdepth3:
toPlot.append(pcl_gt3)
if hasGTdepth4:
toPlot.append(pcl_gt4)
toPlotClear = list(toPlot)
toPlot.append(outlier_cloud1)
toPlot.append(outlier_cloud2)
toPlot.append(outlier_cloud3)
toPlot.append(outlier_cloud4)
o3d.visualization.draw_geometries(toPlot)
o3d.visualization.draw_geometries(toPlotClear)
toPlot = [pcl1, pcl2, pcl3, pcl4]
if hasGTdepth1:
toPlot.append(pcl_gt1)
if hasGTdepth2:
toPlot.append(pcl_gt2)
if hasGTdepth3:
toPlot.append(pcl_gt3)
if hasGTdepth4:
toPlot.append(pcl_gt4)
o3d.visualization.draw_geometries(toPlot)
bbox = o3d.geometry.AxisAlignedBoundingBox(min_bound=(-1000, -1000, -1),
max_bound=(1000, 1000, 5))
toPlot = [
pcl1.crop(bbox),
pcl2.crop(bbox),
pcl3.crop(bbox),
pcl4.crop(bbox)
]
if hasGTdepth1:
toPlot.append(pcl_gt1)
if hasGTdepth2:
toPlot.append(pcl_gt2)
if hasGTdepth3:
toPlot.append(pcl_gt3)
if hasGTdepth4:
toPlot.append(pcl_gt4)
o3d.visualization.draw_geometries(toPlot)
rgb1 = image1[0].permute(1, 2, 0).detach().cpu().numpy() * 255
rgb2 = image2[0].permute(1, 2, 0).detach().cpu().numpy() * 255
rgb3 = image3[0].permute(1, 2, 0).detach().cpu().numpy() * 255
rgb4 = image4[0].permute(1, 2, 0).detach().cpu().numpy() * 255
# Prepare inverse depth
viz_pred_inv_depth1 = viz_inv_depth(pred_inv_depth1[0]) * 255
viz_pred_inv_depth2 = viz_inv_depth(pred_inv_depth2[0]) * 255
viz_pred_inv_depth3 = viz_inv_depth(pred_inv_depth3[0]) * 255
viz_pred_inv_depth4 = viz_inv_depth(pred_inv_depth4[0]) * 255
# Concatenate both vertically
image1 = np.concatenate([rgb1, viz_pred_inv_depth1], 0)
image2 = np.concatenate([rgb2, viz_pred_inv_depth2], 0)
image3 = np.concatenate([rgb3, viz_pred_inv_depth3], 0)
image4 = np.concatenate([rgb4, viz_pred_inv_depth4], 0)
# Save visualization
print('Saving {} to {}'.format(
pcolor(input_file1, 'cyan', attrs=['bold']),
pcolor(output_file1, 'magenta', attrs=['bold'])))
imwrite(output_file1, image1[:, :, ::-1])
print('Saving {} to {}'.format(
pcolor(input_file2, 'cyan', attrs=['bold']),
pcolor(output_file2, 'magenta', attrs=['bold'])))
imwrite(output_file2, image2[:, :, ::-1])
print('Saving {} to {}'.format(
pcolor(input_file3, 'cyan', attrs=['bold']),
pcolor(output_file3, 'magenta', attrs=['bold'])))
imwrite(output_file3, image3[:, :, ::-1])
print('Saving {} to {}'.format(
pcolor(input_file4, 'cyan', attrs=['bold']),
pcolor(output_file4, 'magenta', attrs=['bold'])))
imwrite(output_file4, image4[:, :, ::-1])
def infer_plot_and_save_3D_pcl(input_file, output_file, model_wrapper,
image_shape, half, save):
"""
Process a single input file to produce and save visualization
Parameters
----------
input_file : str
Image file
output_file : str
Output file, or folder where the output will be saved
model_wrapper : nn.Module
Model wrapper used for inference
image_shape : Image shape
Input image shape
half: bool
use half precision (fp16)
save: str
Save format (npz or png)
"""
if not is_image(output_file):
# If not an image, assume it's a folder and append the input name
os.makedirs(output_file, exist_ok=True)
output_file = os.path.join(output_file, os.path.basename(input_file))
# change to half precision for evaluation if requested
dtype = torch.float16 if half else None
# Load image
image = load_image(input_file).convert('RGB')
# Resize and to tensor
image = resize_image(image, image_shape)
image = to_tensor(image).unsqueeze(0)
# Send image to GPU if available
if torch.cuda.is_available():
image = image.to('cuda:{}'.format(rank()), dtype=dtype)
# Depth inference (returns predicted inverse depth)
pred_inv_depth = model_wrapper.depth(image)
pred_depth = inv2depth(pred_inv_depth)
base_folder_str = get_base_folder(input_file)
split_type_str = get_split_type(input_file)
seq_name_str = get_sequence_name(input_file)
camera_str = get_camera_name(input_file)
calib_data = {}
calib_data[camera_str] = read_raw_calib_files_camera_valeo(
base_folder_str, split_type_str, seq_name_str, camera_str)
path_to_theta_lut = get_path_to_theta_lut(input_file)
path_to_ego_mask = get_path_to_ego_mask(input_file)
poly_coeffs, principal_point, scale_factors = get_intrinsics(
input_file, calib_data)
poly_coeffs = torch.from_numpy(poly_coeffs).unsqueeze(0)
principal_point = torch.from_numpy(principal_point).unsqueeze(0)
scale_factors = torch.from_numpy(scale_factors).unsqueeze(0)
pose_matrix = torch.from_numpy(
get_extrinsics_pose_matrix(input_file, calib_data)).unsqueeze(0)
pose_tensor = Pose(pose_matrix)
ego_mask = np.load(path_to_ego_mask)
not_masked = ego_mask.astype(bool).reshape(-1)
cam = CameraFisheye(path_to_theta_lut=[path_to_theta_lut],
path_to_ego_mask=[path_to_ego_mask],
poly_coeffs=poly_coeffs.float(),
principal_point=principal_point.float(),
scale_factors=scale_factors.float(),
Tcw=pose_tensor)
if torch.cuda.is_available():
cam = cam.to('cuda:{}'.format(rank()), dtype=dtype)
world_points = cam.reconstruct(pred_depth, frame='w')
world_points = world_points[0].cpu().numpy()
world_points = world_points.reshape((3, -1)).transpose()
gt_depth_file = get_depth_file(input_file)
gt_depth = np.load(gt_depth_file)['velodyne_depth'].astype(np.float32)
gt_depth = torch.from_numpy(gt_depth).unsqueeze(0).unsqueeze(0)
if torch.cuda.is_available():
gt_depth = gt_depth.to('cuda:{}'.format(rank()), dtype=dtype)
gt_depth_3d = cam.reconstruct(gt_depth, frame='w')
gt_depth_3d = gt_depth_3d[0].cpu().numpy()
gt_depth_3d = gt_depth_3d.reshape((3, -1)).transpose()
world_points = world_points[not_masked]
gt_depth_3d = gt_depth_3d[not_masked]
pcl = o3d.geometry.PointCloud()
pcl.points = o3d.utility.Vector3dVector(world_points)
img_numpy = image[0].cpu().numpy()
img_numpy = img_numpy.reshape((3, -1)).transpose()
img_numpy = img_numpy[not_masked]
pcl.colors = o3d.utility.Vector3dVector(img_numpy)
#pcl.paint_uniform_color([1.0, 0.0, 0])
#print("Radius oulier removal")
#cl, ind = pcl.remove_radius_outlier(nb_points=10, radius=0.5)
#display_inlier_outlier(pcl, ind)
remove_outliers = True
if remove_outliers:
cl, ind = pcl.remove_statistical_outlier(nb_neighbors=10,
std_ratio=1.3)
#display_inlier_outlier(pcl, ind)
inlier_cloud = pcl.select_by_index(ind)
#inlier_cloud.paint_uniform_color([0.0, 1.0, 0])
outlier_cloud = pcl.select_by_index(ind, invert=True)
outlier_cloud.paint_uniform_color([0.0, 0.0, 1.0])
pcl_gt = o3d.geometry.PointCloud()
pcl_gt.points = o3d.utility.Vector3dVector(gt_depth_3d)
pcl_gt.paint_uniform_color([1.0, 0.0, 0])
if remove_outliers:
o3d.visualization.draw_geometries(
[inlier_cloud, pcl_gt, outlier_cloud])
o3d.visualization.draw_geometries([inlier_cloud, pcl_gt])
o3d.visualization.draw_geometries([pcl, pcl_gt])
rgb = image[0].permute(1, 2, 0).detach().cpu().numpy() * 255
# Prepare inverse depth
viz_pred_inv_depth = viz_inv_depth(pred_inv_depth[0]) * 255
# Concatenate both vertically
image = np.concatenate([rgb, viz_pred_inv_depth], 0)
# Save visualization
print('Saving {} to {}'.format(
pcolor(input_file, 'cyan', attrs=['bold']),
pcolor(output_file, 'magenta', attrs=['bold'])))
imwrite(output_file, image[:, :, ::-1])
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 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
