skip empty and to relative coords

mindee · Nov 29, 2021 · acec9df · acec9df
1 parent 578428c
commit acec9df
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 6 deletions.
diff --git a/doctr/datasets/icdar2003.py b/doctr/datasets/icdar2003.py
@@ -81,10 +81,18 @@ def __init__(
                     for rect in rectangles
                 ]
 
-            # Get the labels
-            labels = [lab.text for rect in rectangles for lab in rect if lab.text]
-
-            self.data.append((name.text, dict(boxes=np.asarray(_boxes, dtype=np_dtype), labels=labels)))
+            # filter images without boxes
+            if _boxes:
+                # Convert them to relative
+                w, h = int(resolution.attrib['x']), int(resolution.attrib['y'])
+                boxes = np.asarray(_boxes, dtype=np_dtype)
+                boxes[:, [0, 2]] /= w
+                boxes[:, [1, 3]] /= h
+
+                # Get the labels
+                labels = [lab.text for rect in rectangles for lab in rect if lab.text]
+
+                self.data.append((name.text, dict(boxes=boxes, labels=labels)))
 
         self.root = tmp_root
 

diff --git a/tests/pytorch/test_datasets_pt.py b/tests/pytorch/test_datasets_pt.py
@@ -34,7 +34,7 @@ def test_visiondataset():
         ['SVT', False, [512, 512], 249, False],
         ['SynthText', True, [512, 512], 27, True],  # Actual set has 772875 samples
         ['SynthText', False, [512, 512], 3, False],  # Actual set has 85875 samples
-        ['ICDAR2003', True, [512, 512], 250, True],
+        ['ICDAR2003', True, [512, 512], 246, True],
         ['ICDAR2003', False, [512, 512], 249, False],
     ],
 )

diff --git a/tests/tensorflow/test_datasets_tf.py b/tests/tensorflow/test_datasets_tf.py
@@ -24,7 +24,7 @@
         ['SVT', False, [512, 512], 249, False],
         ['SynthText', True, [512, 512], 27, True],  # Actual set has 772875 samples
         ['SynthText', False, [512, 512], 3, False],  # Actual set has 85875 samples
-        ['ICDAR2003', True, [512, 512], 250, True],
+        ['ICDAR2003', True, [512, 512], 246, True],
         ['ICDAR2003', False, [512, 512], 249, False],
     ],
 )