From 5bd10d73305f9645cfe1bcc5f800b2de25b24bd8 Mon Sep 17 00:00:00 2001
From: Viredery <yiran_ding@outlook.com>
Date: Thu, 19 Jul 2018 13:54:27 +0800
Subject: [PATCH] bug fix: the output_shape of roi_gt_class_ids is incorrect

---
 mrcnn/model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mrcnn/model.py b/mrcnn/model.py
index cc2ee61..68eb5db 100644
--- a/mrcnn/model.py
+++ b/mrcnn/model.py
@@ -670,7 +670,7 @@ class DetectionTargetLayer(KE.Layer):
     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
-- 
GitLab