Spaces:
Runtime error
Runtime error
| #!/usr/bin/env python3 | |
| # -*- coding:utf-8 -*- | |
| # Copyright (c) Megvii, Inc. and its affiliates. | |
| from loguru import logger | |
| from tqdm import tqdm | |
| import torch | |
| 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 | |
| class 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 | |
| ): | |
| """ | |
| 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 | |
| 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] | |
| # preprocessing: resize | |
| 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])] | |
| pred_data = { | |
| "image_id": int(img_id), | |
| "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 import cocoeval as COCOeval | |
| logger.warning("Use standard COCOeval.") | |
| ''' | |
| #from pycocotools.cocoeval import COCOeval | |
| from yolox.layers import COCOeval_opt as 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() | |
| return cocoEval.stats[0], cocoEval.stats[1], info | |
| else: | |
| return 0, 0, info | |