提交 0df74604 编写于 作者: M Megvii Engine Team

fix(data): fix contiguous id

GitOrigin-RevId: 7f79cda0b51f29e8a3ec9678df36fb38aede8ef5
上级 fff2cdc7
......@@ -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 = {
......
......@@ -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 = {
......
......@@ -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],
]
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册