def load_dataset(params):
print("Loading the dataset...")
if params['nyu_dataset']:
dataset = NYUDataset("../data/nyudepthv2/train", split='train')
test_dataset = NYUDataset("../data/nyudepthv2/val", split='val')
else:
dataset = Datasets.FastDepthDataset(
params["training_dataset_paths"],
split='train',
depth_min=params["depth_min"],
depth_max=params["depth_max"],
input_shape_model=(224, 224),
disparity=params["predict_disparity"],
random_crop=params["random_crop"])
test_dataset = Datasets.FastDepthDataset(
params["test_dataset_paths"],
split='val',
depth_min=params["depth_min"],
depth_max=params["depth_max"],
input_shape_model=(224, 224),
disparity=params["predict_disparity"],
random_crop=False)
# Make training/validation split
train_val_split_lengths = utils.get_train_val_split_lengths(
params["train_val_split"], len(dataset))
train_dataset, val_dataset = torch.utils.data.random_split(
dataset, train_val_split_lengths)
params["num_training_examples"] = len(train_dataset)
params["num_validation_examples"] = len(val_dataset)
# DataLoaders
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=params["batch_size"],
shuffle=True,
num_workers=params["num_workers"],
pin_memory=True)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=params["batch_size"],
shuffle=True,
num_workers=params["num_workers"],
pin_memory=True)
test_loader = torch.utils.data.DataLoader(
test_dataset,
batch_size=params["batch_size"],
shuffle=False,
num_workers=params["num_workers"],
pin_memory=True)
return train_loader, val_loader, test_loader
def main(args):
print("Loading config file: ", args.config)
params = utils.load_config_file(args.config)
params["dataset_paths"] = utils.format_dataset_path(
params["dataset_paths"])
if "nyu" not in params:
params["nyu"] = False
# Data loading code
print("Creating data loaders...")
if params["nyu"]:
from dataloaders.nyu import NYUDataset
val_dataset = NYUDataset(params["dataset_paths"], split='val')
else:
val_dataset = Datasets.FastDepthDataset(params["dataset_paths"],
split='val',
depth_min=params["depth_min"],
depth_max=params["depth_max"],
input_shape_model=(224, 224),
random_crop=False)
# set batch size to be 1 for validation
data_loader = torch.utils.data.DataLoader(
val_dataset,
batch_size=1,
shuffle=False,
num_workers=params["num_workers"],
pin_memory=True)
# Set GPU
params["device"] = torch.device(
"cuda:{}".format(params["device"])
if params["device"] >= 0 and torch.cuda.is_available() else "cpu")
print("Using device", params["device"])
print("Loading models...")
models = []
model_names = []
for model_dict in params["models"]:
model_names.append(Path(model_dict["model_path"]).stem)
model, _ = utils.load_model(model_dict, model_dict["model_path"],
params["device"])
model.to(params["device"])
models.append(model)
# Create output directory
output_directory = os.path.join(params["save_folder"],
".".join(model_names))
if not os.path.exists(output_directory):
os.makedirs(output_directory)
params["output_directory"] = output_directory
print("Saving results to " + output_directory)
compare_models(params, data_loader, models)
def main(args):
print("Loading config file: ", args.config)
params = utils.load_config_file(args.config)
params["test_dataset_paths"] = utils.format_dataset_path(
params["test_dataset_paths"])
if args.existing_experiment:
experiment = ExistingExperiment(
api_key="jBFVYFo9VUsy0kb0lioKXfTmM",
previous_experiment=args.existing_experiment)
else:
experiment = Experiment(api_key="jBFVYFo9VUsy0kb0lioKXfTmM",
project_name="fastdepth")
# Data loading code
print("Creating data loaders...")
if args.nyu:
from dataloaders.nyu import NYUDataset
val_dataset = NYUDataset(params["test_dataset_paths"], split='val')
else:
val_dataset = Datasets.FastDepthDataset(params["test_dataset_paths"],
split='val',
depth_min=params["depth_min"],
depth_max=params["depth_max"],
input_shape_model=(224, 224))
# set batch size to be 1 for validation
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=1,
shuffle=True,
num_workers=params["num_workers"],
pin_memory=True)
# Set GPU
params["device"] = torch.device(
"cuda:{}".format(params["device"])
if params["device"] >= 0 and torch.cuda.is_available() else "cpu")
print("Using device", params["device"])
print("Loading model '{}'".format(args.model))
if not args.nyu:
model, _ = utils.load_model(params, args.model, params["device"])
else:
# Maintain compatibility for fastdepth NYU model format
state_dict = torch.load(args.model, map_location=params["device"])
model = models.MobileNetSkipAdd(output_size=(224, 224),
pretrained=True)
model.load_state_dict(state_dict)
params["start_epoch"] = 0
model.to(params["device"])
# Create output directory
output_directory = os.path.join(os.path.dirname(args.model), "images")
if not os.path.exists(output_directory):
os.makedirs(output_directory)
params["experiment_dir"] = output_directory
print("Saving results to " + output_directory)
evaluate(params, val_loader, model, experiment)
