add an attribute for CocoDataset

6d9ee37b · Kai Chen · c96600b1 · 6d9ee37b
Commit 6d9ee37b authored Oct 21, 2018 by Kai Chen
Hide whitespace changes
Inline Side-by-side

Showing with 5 additions and 2 deletions

mmdet/datasets/coco.py mmdet/datasets/coco.py +5 -2

No files found.
--- a/mmdet/datasets/coco.py
+++ b/mmdet/datasets/coco.py
@@ -8,8 +8,11 @@ class CocoDataset(CustomDataset):
    def load_annotations(self, ann_file):
        self.coco = COCO(ann_file)
-        cat_ids = self.coco.getCatIds()
+        self.cat_ids = self.coco.getCatIds()
-        self.cat2label = {cat_id: i + 1 for i, cat_id in enumerate(cat_ids)}
+        self.cat2label = {
+            cat_id: i + 1
+            for i, cat_id in enumerate(self.cat_ids)
+        }
        self.img_ids = self.coco.getImgIds()
        img_infos = []
        for i in self.img_ids: