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)