def get_model(model_name, task):
if task == "seg":
try:
return {
"fcn": fcn_seg(n_classes=19),
"frrnA": frrn_seg(model_type="A", n_classes=19),
"segnet": segnet_seg(n_classes=19),
"deeplab": deeplab_seg(n_classes=19),
"dispnet": dispnet_seg(n_classes=19),
"fcrn": fcrn_seg(n_classes=19),
}[model_name]
except:
raise ("Model {} not available".format(model_name))
elif task == "depth":
try:
return {
"fcn": fcn_depth(),
"frrnA": frrn_depth(model_type="A"),
"segnet": segnet_depth(),
"deeplab": deeplab_depth(),
"dispnet": dispnet_depth(),
"fcrn": fcrn_depth(),
}[model_name]
except:
raise ("Model {} not available".format(model_name))
def main():
img = m.imread(args.img_path).astype(np.float32)
# input image preprocessing, need to match the training settings of the pretrained model
img = imresize(img, (args.height, args.width)).astype(np.float32) # [128, 416, 3]
img = ((img / 255 - 0.5) / 0.5)
img = np.transpose(img, (2, 0, 1))
img = torch.from_numpy(img).unsqueeze(0) # tensor [1, 3, 128, 416]
# load pretrained model
model = dispnet_depth() # TODO: modify model name
weights = torch.load(args.model_path, map_location=lambda storage, loc: storage)
model.load_state_dict(weights['model_state'])
model.eval()
output = model(img).cpu().numpy()[0,0]
y1, y2 = int(0.40810811 * output.shape[0]), int(0.99189189 * output.shape[0])
x1, x2 = int(0 * output.shape[1]), int(1 * output.shape[1])
output_cut = output[y1:y2, x1:x2]
# output_cut = 1/output_cut # TODO: not for dispnet
output_upper = np.full((y1, args.width), np.min(output_cut), dtype=float)
output_cut = (output_cut - np.min(output_cut)) / np.max(output_cut)
output_final = np.concatenate((output_upper, output_cut), axis=0)
m.imsave("output_predict_img/dispnet_output_depth.png", output_final) # for dispnet
def get_model(model_name):
try:
return {
"fcn": fcn_depth(),
"frrnA": frrn_depth(model_type="A"),
"segnet": segnet_depth(),
"deeplab": deeplab_depth(),
"dispnet": dispnet_depth(),
"fcrn": fcrn_depth(),
}[model_name]
except:
raise ("Model {} not available".format(model_name))