fix in collate_fn

odulcy-mindee · Jan 6, 2024 · 1b70baa · 1b70baa
1 parent 7f2600f
commit 1b70baa
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/doctr/datasets/datasets/tensorflow.py b/doctr/datasets/datasets/tensorflow.py
@@ -49,10 +49,18 @@ def _read_sample(self, index: int) -> Tuple[tf.Tensor, Any]:
 
     @staticmethod
     def collate_fn(samples: List[Tuple[tf.Tensor, Any]]) -> Tuple[tf.Tensor, List[Any]]:
-        images, targets = zip(*samples)
+        # FIXME
+        # problems with some shape != 3
+        images, targets = [], []
+        for sample in samples:
+            if sample[0].shape[-1] == 3:
+                images.append(sample[0])
+                targets.append(sample[1])
+
+        # images, targets = zip(*samples)
         images = tf.stack(images, axis=0)
 
-        return images, list(targets)
+        return images, targets
 
 
 class VisionDataset(AbstractDataset, _VisionDataset):  # noqa: D101