fix --verbose error on coco infer

furyhawk · Sep 20, 2022 · 0350f5f · 0350f5f
1 parent 2618785
commit 0350f5f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/yolov6/core/evaler.py b/yolov6/core/evaler.py
@@ -278,7 +278,7 @@ def eval_model(self, pred_results, model, dataloader, task):
                     if ann_i["ignore"]:
                         continue
                     val_dataset_anns_count += 1
-                    nc_i = self.coco80_to_coco91_class().index(ann_i['category_id'])
+                    nc_i = self.coco80_to_coco91_class().index(ann_i['category_id']) if self.is_coco else ann_i['category_id']
                     label_count_dicts[nc_i]["images"].add(ann_i["image_id"])
                     label_count_dicts[nc_i]["anns"] += 1