text: #!/usr/bin/env python3
  # -*- coding:utf-8 -*-
  # Copyright (c) Megvii, Inc. and its affiliates.
  
  from loguru import logger
  from tqdm import tqdm
  
  import torch
  import torchvision
  
  from yolox.utils import (
      gather,
      is_main_process,
      postprocess,
      synchronize,
      time_synchronized,
      xyxy2xywh
  )
  
  import contextlib
  import io
  import itertools
  import json
  import tempfile
  import time
  
  from tabulate import tabulate
  import numpy as np
  from exps.data.argoverse_class import ARGOVERSE_CLASSES
  
  def per_class_mAP_table(coco_eval, class_names=ARGOVERSE_CLASSES, headers=["class", "AP"], colums=2):
      per_class_mAP = {}
      precisions = coco_eval.eval["precision"]
      # precision has dims (iou, recall, cls, area range, max dets)
      assert len(class_names) == precisions.shape[2]
  
      for idx, name in enumerate(class_names):
          # area range index 0: all area ranges
          # max dets index -1: typically 100 per image
          precision = precisions[:, :, idx, 0, -1]
          precision = precision[precision > -1]
          ap = np.mean(precision) if precision.size else float("nan")
          per_class_mAP[name] = float(ap * 100)
  
      num_cols = min(colums, len(per_class_mAP) * len(headers))
      result_pair = [x for pair in per_class_mAP.items() for x in pair]
      row_pair = itertools.zip_longest(*[result_pair[i::num_cols] for i in range(num_cols)])
      table_headers = headers * (num_cols // len(headers))
      table = tabulate(
          row_pair, tablefmt="pipe", floatfmt=".3f", headers=table_headers, numalign="left",
      )
      return table
  
  
  
  class ONEX_COCOEvaluator:
      """
      COCO AP Evaluation class.  All the data in the val2017 dataset are processed
      and evaluated by COCO API.
      """
  
      def __init__(
              self, dataloader, img_size, confthre, nmsthre, num_classes, testdev=False, per_class_mAP=True,
      ):
          """
          Args:
              dataloader (Dataloader): evaluate dataloader.
              img_size (int): image size after preprocess. images are resized
                  to squares whose shape is (img_size, img_size).
              confthre (float): confidence threshold ranging from 0 to 1, which
                  is defined in the config file.
              nmsthre (float): IoU threshold of non-max supression ranging from 0 to 1.
          """
          self.dataloader = dataloader
          self.img_size = img_size
          self.confthre = confthre
          self.nmsthre = nmsthre
          self.num_classes = num_classes
          self.testdev = testdev
  
          self.per_class_mAP = per_class_mAP
  
      def evaluate(
              self,
              model,
              distributed=False,
              half=False,
              trt_file=None,
              decoder=None,
              test_size=None,
      ):
          """
          COCO average precision (AP) Evaluation. Iterate inference on the test dataset
          and the results are evaluated by COCO API.
  
          NOTE: This function will change training mode to False, please save states if needed.
  
          Args:
              model : model to evaluate.
  
          Returns:
              ap50_95 (float) : COCO AP of IoU=50:95
              ap50 (float) : COCO AP of IoU=50
              summary (sr): summary info of evaluation.
          """
          # TODO half to amp_test
          tensor_type = torch.cuda.HalfTensor if half else torch.cuda.FloatTensor
          model = model.eval()
          if half:
              model = model.half()
          ids = []
          data_list = []
          progress_bar = tqdm if is_main_process() else iter
  
          inference_time = 0
          nms_time = 0
          n_samples = len(self.dataloader) - 1
  
          if trt_file is not None:
              from torch2trt import TRTModule
  
              model_trt = TRTModule()
              model_trt.load_state_dict(torch.load(trt_file))
  
              x = torch.ones(1, 3, test_size[0], test_size[1]).cuda()
              model(x)
              model = model_trt
  
          for cur_iter, (imgs, _, info_imgs, ids) in enumerate(
                  progress_bar(self.dataloader)
          ):
              with torch.no_grad():
                  imgs = imgs.type(tensor_type)
  
                  # skip the the last iters since batchsize might be not enough for batch inference
                  is_time_record = cur_iter < len(self.dataloader) - 1
                  if is_time_record:
                      start = time.time()
  
                  outputs = model(imgs)
                  if decoder is not None:
                      outputs = decoder(outputs, dtype=outputs.type())
  
                  if is_time_record:
                      infer_end = time_synchronized()
                      inference_time += infer_end - start
  
                  outputs = postprocess(
                      outputs, self.num_classes, self.confthre, self.nmsthre
                  )
                  if is_time_record:
                      nms_end = time_synchronized()
                      nms_time += nms_end - infer_end
  
              data_list.extend(self.convert_to_coco_format(outputs, info_imgs, ids))
  
          statistics = torch.cuda.FloatTensor([inference_time, nms_time, n_samples])
          if distributed:
              data_list = gather(data_list, dst=0)
              data_list = list(itertools.chain(*data_list))
              torch.distributed.reduce(statistics, dst=0)
  
          eval_results = self.evaluate_prediction(data_list, statistics)
          synchronize()
          return eval_results
  
      def convert_to_coco_format(self, outputs, info_imgs, ids):
          data_list = []
          for (output, img_h, img_w, img_id) in zip(
                  outputs, info_imgs[0], info_imgs[1], ids
          ):
  
              if output is None:
                  continue
              output = output.cpu()
  
              bboxes = output[:, 0:4]
  
              scale = min(
                  self.img_size[0] / float(img_h), self.img_size[1] / float(img_w)
              )
              bboxes /= scale
              bboxes = xyxy2xywh(bboxes)
  
              cls = output[:, 6]
              scores = output[:, 4] * output[:, 5]
  
              for ind in range(bboxes.shape[0]):
                  label = self.dataloader.dataset.class_ids[int(cls[ind])]
  
                  if int(img_id) in [15060,15061]:
                      continue
                  elif self.dataloader.dataset.coco.dataset['images'][int(img_id+1)]['fid'] == 0:
                      continue
                  elif self.dataloader.dataset.coco.dataset['images'][int(img_id)]['fid'] == 0:
                      idd = int(img_id)
                  else:
                      idd = int(img_id + 1)
  
                      pred_data = {
                          "image_id": idd,  
                          "category_id": label,
                          "bbox": bboxes[ind].numpy().tolist(),
                          "score": scores[ind].numpy().item(),
                          "segmentation": [],
                      }  # COCO json format
                      data_list.append(pred_data)
                      
          return data_list
  
      def evaluate_prediction(self, data_dict, statistics):
          if not is_main_process():
              return 0, 0, None
  
          logger.info("Evaluate in main process...")
  
          annType = ["segm", "bbox", "keypoints"]
  
          inference_time = statistics[0].item()
          nms_time = statistics[1].item()
          n_samples = statistics[2].item()
  
          a_infer_time = 1000 * inference_time / (n_samples * self.dataloader.batch_size)
          a_nms_time = 1000 * nms_time / (n_samples * self.dataloader.batch_size)
  
          time_info = ", ".join(
              [
                  "Average {} time: {:.2f} ms".format(k, v)
                  for k, v in zip(
                  ["forward", "NMS", "inference"],
                  [a_infer_time, a_nms_time, (a_infer_time + a_nms_time)],
              )
              ]
          )
  
          info = time_info + "\n"
  
          # Evaluate the Dt (detection) json comparing with the ground truth
          if len(data_dict) > 0:
              cocoGt = self.dataloader.dataset.coco
  
              # TODO: since pycocotools can't process dict in py36, write data to json file.
              if self.testdev:
                  json.dump(data_dict, open("./yolox_testdev_2017.json", "w"))
                  cocoDt = cocoGt.loadRes("./yolox_testdev_2017.json")
              else:
                  _, tmp = tempfile.mkstemp()
                  json.dump(data_dict, open(tmp, "w"))
                  cocoDt = cocoGt.loadRes(tmp)
                  
              try:
                  from yolox.layers import COCOeval_opt as COCOeval
              except ImportError:
                  from pycocotools.cocoeval import COCOeval
  
  
              cocoEval = COCOeval(cocoGt, cocoDt, annType[1])
              cocoEval.evaluate()
              cocoEval.accumulate()
              redirect_string = io.StringIO()
              with contextlib.redirect_stdout(redirect_string):
                  cocoEval.summarize()
              info += redirect_string.getvalue()
              if self.per_class_mAP:
                  info += "per class mAP:\n" + per_class_mAP_table(cocoEval)
              return cocoEval.stats[0], cocoEval.stats[1], info
          else:
              return 0, 0, info
  