Skip to content

GitLab

Commit 6d9ee37b authored by Kai Chen's avatar Kai Chen
Browse files

add an attribute for CocoDataset

parent c96600b1
Hide whitespace changes
Inline Side-by-side
Showing with 5 additions and 2 deletions
mmdet/datasets/coco.py
View file @ 6d9ee37b
...@@ -8,8 +8,11 @@ class CocoDataset(CustomDataset): ...@@ -8,8 +8,11 @@ class CocoDataset(CustomDataset):
def load_annotations(self, ann_file): def load_annotations(self, ann_file):
self.coco = COCO(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() self.img_ids = self.coco.getImgIds()
img_infos = [] img_infos = []
for i in self.img_ids: for i in self.img_ids:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment