def __init__(self, dataset='nyu', device='cuda:0'):
self.toTensor = ToTensor()
self.device = device
if dataset == 'nyu':
self.min_depth = 1e-3
self.max_depth = 20
self.saving_factor = 1000 # used to save in 16 bit
model = UnetAdaptiveBins.build(n_bins=256,
min_val=self.min_depth,
max_val=self.max_depth)
pretrained_path = "./pretrained/AdaBins_nyu.pt"
elif dataset == 'kitti':
self.min_depth = 1e-3
self.max_depth = 80
self.saving_factor = 256
model = UnetAdaptiveBins.build(n_bins=256,
min_val=self.min_depth,
max_val=self.max_depth)
pretrained_path = "./pretrained/AdaBins_kitti.pt"
else:
raise ValueError(
"dataset can be either 'nyu' or 'kitti' but got {}".format(
dataset))
model, _, _ = model_io.load_checkpoint(pretrained_path, model)
model.eval()
self.model = model.to(self.device)
default=10)
parser.add_argument('--eigen_crop',
help='if set, crops according to Eigen NIPS14',
action='store_true')
parser.add_argument('--garg_crop',
help='if set, crops according to Garg ECCV16',
action='store_true')
parser.add_argument('--do_kb_crop',
help='Use kitti benchmark cropping',
action='store_true')
if sys.argv.__len__() == 2:
arg_filename_with_prefix = '@' + sys.argv[1]
args = parser.parse_args([arg_filename_with_prefix])
else:
args = parser.parse_args()
# args = parser.parse_args()
args.gpu = int(args.gpu) if args.gpu is not None else 0
args.distributed = False
device = torch.device('cuda:{}'.format(args.gpu))
test = DepthDataLoader(args, 'online_eval').data
model = UnetAdaptiveBins.build(n_bins=args.n_bins,
min_val=args.min_depth,
max_val=args.max_depth,
norm='linear').to(device)
model = model_io.load_checkpoint(args.checkpoint_path, model)[0]
model = model.eval()
eval(model, test, args, gpus=[device])