Merge pull request #60 from avilash/master

Removing hard coded number of classes for VOC dataset
daijifeng001 · Jun 9, 2017 · e6be483 · e6be483
2 parents e31c98e + 94649d9
commit e6be483
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/lib/pylayer/stage_bridge_layer.py b/lib/pylayer/stage_bridge_layer.py
@@ -49,11 +49,11 @@ def setup(self, bottom, top):
             self._top_name_map['mask_weight'] = 3
             top[4].reshape(1, 4)
             self._top_name_map['gt_mask_info'] = 4
-            top[5].reshape(1, 21 * 4)
+            top[5].reshape(1, self._num_classes * 4)
             self._top_name_map['bbox_targets'] = 5
-            top[6].reshape(1, 21 * 4)
+            top[6].reshape(1, self._num_classes * 4)
             self._top_name_map['bbox_inside_weights'] = 6
-            top[7].reshape(1, 21 * 4)
+            top[7].reshape(1, self._num_classes * 4)
             self._top_name_map['bbox_outside_weights'] = 7
         elif self._phase == 'TEST':
             top[0].reshape(1, 5)