fix mask_to_seq bug

Yoctol · Jun 19, 2017 · 3d998ae · 3d998ae
1 parent dbcd736
commit 3d998ae
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 4 deletions.
diff --git a/test/wrapper/test_mask_to_seq.py b/test/wrapper/test_mask_to_seq.py
@@ -42,13 +42,24 @@ def test_output_shape(self):
     def test_image_data_mask_value(self):
         result = self.model.predict(self.data)
         np.testing.assert_almost_equal(
-            result,
+            result[:, :self.x_start, :],
             np.zeros((
                 self.batch_size,
-                self.x,
+                self.x_start,
                 self.y * self.channel_size
             ))
         )
+        np.testing.assert_almost_equal(
+            result[:, self.x_end:, :],
+            np.zeros((
+                self.batch_size,
+                self.x - self.x_end,
+                self.y * self.channel_size
+            ))
+        )
+        self.assertTrue(
+            np.sum(result[:, self.x_start:self.x_end, :], dtype=bool)
+        )
 
     def test_seq_data_mask_value(self):
         result = self.model.predict(self.seq_data)
@@ -71,7 +82,21 @@ def test_image_data_mask(self):
             session=K.get_session(),
             feed_dict={self.model.input: self.data}
         )
-        self.assertFalse(np.any(mask))
+        self.assertTrue(
+            np.all(
+                mask[:, self.x_start:self.x_end]
+            )
+        )
+        self.assertFalse(
+            np.any(
+                mask[:, :self.x_start]
+            )
+        )
+        self.assertFalse(
+            np.any(
+                mask[:, self.x_end:]
+            )
+        )
 
     def test_seq_data_mask(self):
         mask_cache_key = str(id(self.model.input)) + '_' + str(id(None))

diff --git a/yklz/wrapper/mask_to_seq.py b/yklz/wrapper/mask_to_seq.py
@@ -27,7 +27,7 @@ def compute_mask(self, inputs, mask=None):
 
         reduce_time = len(mask_shape) - 2
         for _ in range(reduce_time):
-            mask_tensor = K.all(mask_tensor, -1)
+            mask_tensor = K.any(mask_tensor, -1)
         return mask_tensor
 
     def call(self, inputs, mask=None):