chore(evaluator): support custom dataset (#1131)

chore(evaluator): support custom dataset Co-authored-by: Feng Wang <wangfeng19950315@163.com>
Megvii-BaseDetection · Feb 21, 2022 · 2eda8fe · 2eda8fe
1 parent 57309c7
commit 2eda8fe
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/yolox/evaluators/coco_evaluator.py b/yolox/evaluators/coco_evaluator.py
@@ -280,10 +280,14 @@ def evaluate_prediction(self, data_dict, statistics):
             with contextlib.redirect_stdout(redirect_string):
                 cocoEval.summarize()
             info += redirect_string.getvalue()
+            cat_ids = list(cocoGt.cats.keys())
+            cat_names = [cocoGt.cats[catId]['name'] for catId in sorted(cat_ids)]
             if self.per_class_AP:
-                info += "per class AP:\n" + per_class_AP_table(cocoEval) + "\n"
+                AP_table = per_class_AP_table(cocoEval, class_names=cat_names)
+                info += "per class AP:\n" + AP_table + "\n"
             if self.per_class_AR:
-                info += "per class AR:\n" + per_class_AR_table(cocoEval) + "\n"
+                AR_table = per_class_AR_table(cocoEval, class_names=cat_names)
+                info += "per class AR:\n" + AR_table + "\n"
             return cocoEval.stats[0], cocoEval.stats[1], info
         else:
             return 0, 0, info