From 0df74604bd6b8b39c7fc8f2a55c640c10d3c7f48 Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Thu, 16 Jul 2020 17:09:53 +0800
Subject: [PATCH] fix(data): fix contiguous id

GitOrigin-RevId: 7f79cda0b51f29e8a3ec9678df36fb38aede8ef5
---
 .../megengine/data/dataset/vision/coco.py          |  2 +-
 .../megengine/data/dataset/vision/objects365.py    |  2 +-
 python_module/megengine/data/dataset/vision/voc.py | 40 ++++++++--------------
 3 files changed, 17 insertions(+), 27 deletions(-)

diff --git a/python_module/megengine/data/dataset/vision/coco.py b/python_module/megengine/data/dataset/vision/coco.py
index d247e52b..11366de0 100644
--- a/python_module/megengine/data/dataset/vision/coco.py
+++ b/python_module/megengine/data/dataset/vision/coco.py
@@ -118,7 +118,7 @@ class COCO(VisionDataset):
             self.ids = ids
 
         self.json_category_id_to_contiguous_id = {
-            v: i + 1 for i, v in enumerate(self.cats.keys())
+            v: i + 1 for i, v in enumerate(sorted(self.cats.keys()))
         }
 
         self.contiguous_category_id_to_json_id = {
diff --git a/python_module/megengine/data/dataset/vision/objects365.py b/python_module/megengine/data/dataset/vision/objects365.py
index 7c1481ba..e56e6462 100644
--- a/python_module/megengine/data/dataset/vision/objects365.py
+++ b/python_module/megengine/data/dataset/vision/objects365.py
@@ -81,7 +81,7 @@ class Objects365(VisionDataset):
             self.ids = ids
 
         self.json_category_id_to_contiguous_id = {
-            v: i + 1 for i, v in enumerate(self.cats.keys())
+            v: i + 1 for i, v in enumerate(sorted(self.cats.keys()))
         }
 
         self.contiguous_category_id_to_json_id = {
diff --git a/python_module/megengine/data/dataset/vision/voc.py b/python_module/megengine/data/dataset/vision/voc.py
index 42bf712d..b22fd2fa 100644
--- a/python_module/megengine/data/dataset/vision/voc.py
+++ b/python_module/megengine/data/dataset/vision/voc.py
@@ -75,6 +75,8 @@ class PascalVOC(VisionDataset):
         else:
             raise NotImplementedError
 
+        self.img_infos = dict()
+
     def __getitem__(self, index):
         target = []
         for k in self.order:
@@ -107,9 +109,8 @@ class PascalVOC(VisionDataset):
                 mask = mask[:, :, np.newaxis]
                 target.append(mask)
             elif k == "info":
-                if image is None:
-                    image = cv2.imread(self.images[index], cv2.IMREAD_COLOR)
-                info = [image.shape[0], image.shape[1], self.file_names[index]]
+                info = self.get_img_info(index, image)
+                info = [info["height"], info["width"], info["file_name"]]
                 target.append(info)
             else:
                 raise NotImplementedError
@@ -119,6 +120,17 @@ class PascalVOC(VisionDataset):
     def __len__(self):
         return len(self.images)
 
+    def get_img_info(self, index, image=None):
+        if index not in self.img_infos:
+            if image is None:
+                image = cv2.imread(self.images[index], cv2.IMREAD_COLOR)
+            self.img_infos[index] = dict(
+                height=image.shape[0],
+                width=image.shape[1],
+                file_name=self.file_names[index],
+            )
+        return self.img_infos[index]
+
     def _trans_mask(self, mask):
         label = np.ones(mask.shape[:2]) * 255
         for i in range(len(self.class_colors)):
@@ -171,25 +183,3 @@ class PascalVOC(VisionDataset):
         "train",
         "tvmonitor",
     )
-    class_colors = [
-        [0, 0, 128],
-        [0, 128, 0],
-        [0, 128, 128],
-        [128, 0, 0],
-        [128, 0, 128],
-        [128, 128, 0],
-        [128, 128, 128],
-        [0, 0, 64],
-        [0, 0, 192],
-        [0, 128, 64],
-        [0, 128, 192],
-        [128, 0, 64],
-        [128, 0, 192],
-        [128, 128, 64],
-        [128, 128, 192],
-        [0, 64, 0],
-        [0, 64, 128],
-        [0, 192, 0],
-        [0, 192, 128],
-        [128, 64, 0],
-    ]