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))