Merge pull request #131 from tdiekel/master

Support for training custom datasets
mahyarnajibi · Oct 8, 2019 · 0117d7a · 0117d7a
2 parents 3b82040 + ce7b3bc
commit 0117d7a
Show file tree

Hide file tree

Showing 5 changed files with 9 additions and 9 deletions.
diff --git a/lib/dataset/imdb.py b/lib/dataset/imdb.py
@@ -152,7 +152,7 @@ def create_roidb_from_box_list(self, box_list, mapping_list, gt_roidb):
         self.num_images = len(gt_roidb)
         assert len(box_list) == self.num_images, 'number of boxes matrix must match number of images'
         roidb = []
-        stats = np.zeros(81)
+        stats = np.zeros(self.num_classes)
         for i in range(self.num_images):
             roi_rec = dict()
             roi_rec['image'] = gt_roidb[i]['image']

diff --git a/symbols/faster/mobilenetv2_e2e.py b/symbols/faster/mobilenetv2_e2e.py
@@ -274,7 +274,7 @@ def get_symbol_rcnn(self, cfg, is_train=True):
             fc_new_2_relu = mx.sym.Activation(data=fc_new_2, act_type='relu', name='fc_new_2_relu')
 
             num_reg_classes = 1
-            num_classes = 81
+            num_classes = cfg.dataset.NUM_CLASSES
             cls_score = mx.sym.FullyConnected(name='cls_score', data=fc_new_2_relu, num_hidden=num_classes)
             bbox_pred = mx.sym.FullyConnected(name='bbox_pred', data=fc_new_2_relu, num_hidden=num_reg_classes * 4)
 
@@ -351,7 +351,7 @@ def get_symbol_rcnn(self, cfg, is_train=True):
             fc_new_2_relu = mx.sym.Activation(data=fc_new_2, act_type='relu', name='fc_new_2_relu')
 
             num_reg_classes = 1
-            num_classes = 81
+            num_classes = cfg.dataset.NUM_CLASSES
             cls_score = mx.sym.FullyConnected(name='cls_score', data=fc_new_2_relu, num_hidden=num_classes)
             bbox_pred = mx.sym.FullyConnected(name='bbox_pred', data=fc_new_2_relu, num_hidden=num_reg_classes * 4)
 

diff --git a/symbols/faster/resnet_mx_101_e2e.py b/symbols/faster/resnet_mx_101_e2e.py
@@ -284,7 +284,7 @@ def get_symbol_rcnn(self, cfg, is_train=True):
 
             fc_new_2 = mx.sym.FullyConnected(name='fc_new_2', data=fc_new_1_relu, num_hidden=1024)
             fc_new_2_relu = mx.sym.Activation(data=fc_new_2, act_type='relu', name='fc_new_2_relu')
-            num_classes = 81
+            num_classes = cfg.dataset.NUM_CLASSES
             num_reg_classes = 1
             cls_score = mx.sym.FullyConnected(name='cls_score', data=fc_new_2_relu, num_hidden=num_classes)
             bbox_pred = mx.sym.FullyConnected(name='bbox_pred', data=fc_new_2_relu, num_hidden=num_reg_classes * 4)            
@@ -353,7 +353,7 @@ def get_symbol_rcnn(self, cfg, is_train=True):
 
             fc_new_2 = mx.sym.FullyConnected(name='fc_new_2', data=fc_new_1_relu, num_hidden=1024)
             fc_new_2_relu = mx.sym.Activation(data=fc_new_2, act_type='relu', name='fc_new_2_relu')
-            num_classes = 81
+            num_classes = cfg.dataset.NUM_CLASSES
             num_reg_classes = 1
             cls_score = mx.sym.FullyConnected(name='cls_score', data=fc_new_2_relu, num_hidden=num_classes)
             bbox_pred = mx.sym.FullyConnected(name='bbox_pred', data=fc_new_2_relu, num_hidden=num_reg_classes * 4)

diff --git a/symbols/faster/resnet_mx_101_e2e_mask.py b/symbols/faster/resnet_mx_101_e2e_mask.py
@@ -322,7 +322,7 @@ def get_symbol_rcnn(self, cfg, is_train=True):
 
             fc_new_2 = mx.sym.FullyConnected(name='fc_new_2', data=fc_new_1_relu, num_hidden=1024)
             fc_new_2_relu = mx.sym.Activation(data=fc_new_2, act_type='relu', name='fc_new_2_relu')
-            num_classes = 81
+            num_classes = cfg.dataset.NUM_CLASSES
             num_reg_classes = 1
             cls_score = mx.sym.FullyConnected(name='cls_score', data=fc_new_2_relu, num_hidden=num_classes)
             bbox_pred = mx.sym.FullyConnected(name='bbox_pred', data=fc_new_2_relu, num_hidden=num_reg_classes * 4)
@@ -428,7 +428,7 @@ def get_symbol_rcnn(self, cfg, is_train=True):
 
             fc_new_2 = mx.sym.FullyConnected(name='fc_new_2', data=fc_new_1_relu, num_hidden=1024)
             fc_new_2_relu = mx.sym.Activation(data=fc_new_2, act_type='relu', name='fc_new_2_relu')
-            num_classes = 81
+            num_classes = cfg.dataset.NUM_CLASSES
             num_reg_classes = 1
             cls_score = mx.sym.FullyConnected(name='cls_score', data=fc_new_2_relu, num_hidden=num_classes)
             bbox_pred = mx.sym.FullyConnected(name='bbox_pred', data=fc_new_2_relu, num_hidden=num_reg_classes * 4)

diff --git a/symbols/faster/resnet_mx_50_e2e.py b/symbols/faster/resnet_mx_50_e2e.py
@@ -285,7 +285,7 @@ def get_symbol_rcnn(self, cfg, is_train=True):
 
             fc_new_2 = mx.sym.FullyConnected(name='fc_new_2', data=fc_new_1_relu, num_hidden=1024)
             fc_new_2_relu = mx.sym.Activation(data=fc_new_2, act_type='relu', name='fc_new_2_relu')
-            num_classes = 81
+            num_classes = cfg.dataset.NUM_CLASSES
             num_reg_classes = 1
             cls_score = mx.sym.FullyConnected(name='cls_score', data=fc_new_2_relu, num_hidden=num_classes)
             bbox_pred = mx.sym.FullyConnected(name='bbox_pred', data=fc_new_2_relu, num_hidden=num_reg_classes * 4)            
@@ -354,7 +354,7 @@ def get_symbol_rcnn(self, cfg, is_train=True):
 
             fc_new_2 = mx.sym.FullyConnected(name='fc_new_2', data=fc_new_1_relu, num_hidden=1024)
             fc_new_2_relu = mx.sym.Activation(data=fc_new_2, act_type='relu', name='fc_new_2_relu')
-            num_classes = 81
+            num_classes = cfg.dataset.NUM_CLASSES
             num_reg_classes = 1
             cls_score = mx.sym.FullyConnected(name='cls_score', data=fc_new_2_relu, num_hidden=num_classes)
             bbox_pred = mx.sym.FullyConnected(name='bbox_pred', data=fc_new_2_relu, num_hidden=num_reg_classes * 4)