text: #!/usr/bin/env python
  # -*- coding: utf-8 -*-
  import torch
  from models import YoloNet
  from data_loader import DataTransformBase
  from models import DetectHandler
  
  
  def test_one_image(
      args, total_config, dataset_class, dataset_params,
  ):
      import cv2
      import numpy as np
  
      model_path = args.snapshot
      dataset = args.dataset
      dataset_params = total_config.DATASET_PARAMS[dataset]
      input_size = (dataset_params["img_h"], dataset_params["img_w"])
  
      dataset_instance = dataset_class(data_path=total_config.DATA_PATH)
      num_classes = dataset_instance.num_classes
      model = YoloNet(dataset_config=dataset_params)
      model.load_state_dict(torch.load(model_path)["state_dict"])
      model.eval()
  
      img = cv2.imread(args.image_path)
      orig_img = np.copy(img)
  
      ori_h, ori_w = img.shape[:2]
      h_ratio = ori_h / dataset_params["img_h"]
      w_ratio = ori_w / dataset_params["img_w"]
  
      img = cv2.resize(img, input_size)
  
      img = img / 255.0
      input_x = torch.tensor(img.transpose(2, 0, 1)[np.newaxis, :]).float()
  
      predictions = model(input_x)
  
      detect_handler = DetectHandler(
          num_classes=dataset_params["num_classes"],
          conf_thresh=args.conf_thresh,
          nms_thresh=args.nms_thresh,
          h_ratio=h_ratio,
          w_ratio=w_ratio,
      )
  
      bboxes, scores, classes = detect_handler(predictions)
  
      result = dataset_class.visualize_one_image_util(
          orig_img, dataset_instance.classes, dataset_instance.colors, bboxes, classes,
      )
  
      return orig_img
  
  
  def main(args):
      import cv2
      import os
      from config import Config
  
      total_config = Config()
      if not args.dataset or args.dataset not in total_config.DATASETS.keys():
          raise Exception("specify one of the datasets to use in {}".format(list(total_config.DATASETS.keys())))
      if not args.snapshot or not os.path.isfile(args.snapshot):
          raise Exception("invalid snapshot")
      if not args.image_path or not os.path.isfile(args.image_path):
          raise Exception("invalid image path")
  
      dataset = args.dataset
      dataset_class = total_config.DATASETS[dataset]
      dataset_params = total_config.DATASET_PARAMS[dataset]
  
      result = test_one_image(args, total_config, dataset_class, dataset_params,)
  
      cv2.imshow("result", result)
      cv2.waitKey(0)
      cv2.destroyAllWindows()
      cv2.imwrite("result.jpg", result)
  
  
  if __name__ == "__main__":
      import argparse
  
      parser = argparse.ArgumentParser()
      parser.add_argument("--snapshot", required=True, type=str)
      parser.add_argument(
          "--dataset", type=str, required=True, help="name of the dataset to use",
      )
      parser.add_argument("--image_path", required=True, type=str, help="path to the test image")
      parser.add_argument("--conf_thresh", type=float, default=0.1)
      parser.add_argument("--nms_thresh", type=float, default=0.5)
      parsed_args = parser.parse_args()
  
      main(parsed_args)
  