text: # test end to end benchmark data test
  import argparse
  import os
  
  import cv2
  # import PIL
  import matplotlib.pyplot as plt
  import numpy as np
  import torch
  # import torch.nn as nn
  import torch.nn.functional as F
  
  from models import get_model
  # from utils import convert_state_dict
  
  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  
  
  def unwarp(img, bm):
      w, h = img.shape[0], img.shape[1]
      bm = bm.transpose(1, 2).transpose(2, 3).detach().cpu().numpy()[0, :, :, :]
      bm0 = cv2.blur(bm[:, :, 0], (9, 9))
      bm1 = cv2.blur(bm[:, :, 1], (9, 9))
      bm0 = cv2.resize(bm0, (h, w))
      bm1 = cv2.resize(bm1, (h, w))
      bm = np.stack([bm0, bm1], axis=-1)
      bm = np.expand_dims(bm, 0)
      bm = torch.from_numpy(bm).double()
  
      img = img.astype(float) / 255.0
      img = img.transpose((2, 0, 1))
      img = np.expand_dims(img, 0)
      img = torch.from_numpy(img).double()
  
      res = F.grid_sample(input=img, grid=bm)
      res = res[0].numpy().transpose((1, 2, 0))
  
      return res
  
  
  def test(args, img_path, fname):
      bm_n_classes = 2
  
      bm_img_size = (128, 128)
  
      # Setup image
      print("Read Input Image from : {}".format(img_path))
  
      imgorg = cv2.imread(img_path)
      imgorg = cv2.cvtColor(imgorg, cv2.COLOR_BGR2RGB)
      img = cv2.resize(imgorg, bm_img_size)
      # img = img[:, :, ::-1]
      img = img.astype(float) / 255.0
      img = img.transpose(2, 0, 1)  # NHWC -> NCHW
      img = np.expand_dims(img, 0)
      img = torch.from_numpy(img).float()
  
      img_bm_model = get_model('unet', bm_n_classes, in_channels=3)
      img_bm_model_state = torch.load(args.model_path, map_location=device)['model_state']
      img_bm_model.load_state_dict(img_bm_model_state)
      img_bm_model.eval()
  
      img_bm_model = img_bm_model.to(device)
      images = img.to(device)
  
      with torch.no_grad():
          outputs_bm = img_bm_model(images)
  
          # call unwarp
          uwpred = unwarp(imgorg, outputs_bm)
  
          if args.show:
              f1, axarr1 = plt.subplots(1, 2)
              axarr1[0].imshow(imgorg)
              axarr1[1].imshow(uwpred)
              plt.savefig(dpi=300)
  
          # Save the output
          outp = os.path.join(args.out_path, fname)
          cv2.imwrite(outp, uwpred[:, :, ::-1]*255)
  
  
  if __name__ == '__main__':
      parser = argparse.ArgumentParser(description='Params')
      parser.add_argument('--model_path', type=str,
                          default='experiments/unet_img_bm_docunet_loss/checkpoint/unet_19_0.03266928295688654_best_model.pth',
                          help='Path to the saved image bm model')
      parser.add_argument('--img_path', type=str, default='./eval/input/',
                          help='Path of the input image')
      parser.add_argument('--out_path', type=str, default='./eval/unwarp/',
                          help='Path of the output unwarped image')
      parser.add_argument('--show', dest='show', action='store_true',
                          help='Show the input image and output unwarped')
      parser.set_defaults(show=False)
      args = parser.parse_args()
      if not os.path.exists(args.out_path):
          os.makedirs(args.out_path)
      for fname in os.listdir(args.img_path):
          if '.jpg' in fname or '.JPG' in fname or '.png' in fname:
              img_path = os.path.join(args.img_path, fname)
              test(args, img_path, fname)
  
  
  # python infer.py --model_path ./eval/models/unetnc_doc3d.pkl --show
  