bug fix: the output_shape of roi_gt_class_ids is incorrect

CRAWlab · Sep 21, 2018 · 5bd10d7 · 5bd10d7
1 parent 744fbc2
commit 5bd10d7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/mrcnn/model.py b/mrcnn/model.py
@@ -670,7 +670,7 @@ def call(self, inputs):
     def compute_output_shape(self, input_shape):
         return [
             (None, self.config.TRAIN_ROIS_PER_IMAGE, 4),  # rois
-            (None, 1),  # class_ids
+            (None, self.config.TRAIN_ROIS_PER_IMAGE),  # class_ids
             (None, self.config.TRAIN_ROIS_PER_IMAGE, 4),  # deltas
             (None, self.config.TRAIN_ROIS_PER_IMAGE, self.config.MASK_SHAPE[0],
              self.config.MASK_SHAPE[1])  # masks