Add type to index mapping. Ref #4

iMoonLab · Apr 18, 2019 · e0f06c7 · e0f06c7
1 parent 437e20d
commit e0f06c7
Showing 1 changed file with 12 additions and 2 deletions.
diff --git a/data/ModelNet40.py b/data/ModelNet40.py
@@ -3,6 +3,17 @@
 import torch
 import torch.utils.data as data
 
+type_to_index_map = {
+    'night_stand': 0, 'range_hood': 1, 'plant': 2, 'chair': 3, 'tent': 4,
+    'curtain': 5, 'piano': 6, 'dresser': 7, 'desk': 8, 'bed': 9,
+    'sink': 10,  'laptop':11, 'flower_pot': 12, 'car': 13, 'stool': 14,
+    'vase': 15, 'monitor': 16, 'airplane': 17, 'stairs': 18, 'glass_box': 19,
+    'bottle': 20, 'guitar': 21, 'cone': 22,  'toilet': 23, 'bathtub': 24,
+    'wardrobe': 25, 'radio': 26,  'person': 27, 'xbox': 28, 'bowl': 29,
+    'cup': 30, 'door': 31,  'tv_stand': 32,  'mantel': 33, 'sofa': 34,
+    'keyboard': 35, 'bookshelf': 36,  'bench': 37, 'table': 38, 'lamp': 39
+}
+
 
 class ModelNet40(data.Dataset):
 
@@ -12,13 +23,12 @@ def __init__(self, cfg, part='train'):
         self.part = part
 
         self.data = []
-        type_index = 0
         for type in os.listdir(self.root):
+            type_index = type_to_index_map[type]
             type_root = os.path.join(os.path.join(self.root, type), part)
             for filename in os.listdir(type_root):
                 if filename.endswith('.npz'):
                     self.data.append((os.path.join(type_root, filename), type_index))
-            type_index += 1
 
     def __getitem__(self, i):
         path, type = self.data[i]