update tests

mindee · fg-mindee · Dec 10, 2021 · Nov 13, 2021 · Nov 15, 2021 · Nov 15, 2021
commit ebb7a827610abb7b1b6352327dd92e621d71bfda
diff --git a/doctr/datasets/ic13.py b/doctr/datasets/ic13.py
@@ -4,8 +4,8 @@
 # See LICENSE or go to <https://www.apache.org/licenses/LICENSE-2.0.txt> for full license details.
 
 import os
-from typing import Any, Callable, List, Optional, Tuple, Dict
 from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import numpy as np
 
@@ -15,18 +15,19 @@
 
 
 class IC13(AbstractDataset):
-    """IC13 dataset from
-    `"ICDAR 2013 Robust Reading Competition" <https://rrc.cvc.uab.es/>`_.
+    """IC13 dataset from `"ICDAR 2013 Robust Reading Competition" <https://rrc.cvc.uab.es/>`_.
     Example::
         >>> # NOTE: You need to download both image and label parts from Focused Scene Text challenge Task2.1 2013-2015.
-        >>> # structure: - img_folder: all images (part1 and part2) | - label_folder: all labels (part1 and part2)
         >>> from doctr.datasets import IC13
-        >>> train_set = IC13(img_folder="/path/to/images", label_folder="/path/to/labels", train=True)
+        >>> train_set = IC13(img_folder="/path/to/Challenge2_Training_Task12_Images",
+        >>>                  label_folder="/path/to/Challenge2_Training_Task1_GT")
         >>> img, target = train_set[0]
+        >>> test_set = IC13(img_folder="/path/to/Challenge2_Test_Task12_Images",
+        >>>                 label_folder="/path/to/Challenge2_Test_Task1_GT")
+        >>> img, target = test_set[0]
     Args:
         img_folder: folder with all the images of the dataset
         label_folder: folder with all annotation files for the images
-        train: whether the subset should be the training one
         sample_transforms: composable transformations that will be applied to each image
         rotated_bbox: whether polygons should be considered as rotated bounding box (instead of straight ones)
     """
@@ -35,7 +36,6 @@ def __init__(
         self,
         img_folder: str,
         label_folder: str,
-        train: bool = True,
         sample_transforms: Optional[Callable[[Any], Any]] = None,
         rotated_bbox: bool = False,
     ) -> None:
@@ -51,30 +51,25 @@ def __init__(
         np_dtype = np.float32
 
         img_names = os.listdir(img_folder)
-        split = int(len(img_names) * 0.9)
-        img_names = img_names[:split] if train else img_names[split:]
 
         for img_name in img_names:
 
             img_path = Path(img_folder, img_name)
-            label_path = Path(label_folder, Path(img_name).stem + ".txt")
+            label_path = Path(label_folder, "gt_" + Path(img_name).stem + ".txt")
 
             with open(label_path, "r") as f:
-                _lines = [[x for x in line.split(',')] for line in f.readlines()]
+                _lines = [[x.replace("\"", "").replace(",", "") for x in line.split()] for line in f.readlines()]
                 labels = [line[4].strip() for line in _lines]
-                coords = np.array([np.array(list(map(int, line[:8]))).reshape((4, 2))
-                                  for line in _lines], dtype=np_dtype)
+                # xmin, ymin, xmax, ymax
+                box_targets = np.array([list(map(int, line[:4])) for line in _lines], dtype=np_dtype)
                 if rotated_bbox:
-                    # x_center, y_center, w, h, alpha = 0
-                    mins = coords.min(axis=1)
-                    maxs = coords.max(axis=1)
-                    box_targets = np.concatenate(
-                        ((mins + maxs) / 2, maxs - mins, np.zeros((coords.shape[0], 1))), axis=1)
-                else:
-                    # xmin, ymin, xmax, ymax
-                    box_targets = np.concatenate((coords.min(axis=1), coords.max(axis=1)), axis=1)
-
-                self.data.append((img_path, dict(boxes=box_targets, languages=vocab_lang, labels=labels)))
+                    # x_center, y_center, width, height, 0
+                    box_targets = np.array([[coords[0] + (coords[2] - coords[0]) / 2,
+                                             coords[1] + (coords[3] - coords[1]) / 2,
+                                             (coords[2] - coords[0]),
+                                             (coords[3] - coords[1]), 0.0] for coords in box_targets], dtype=np_dtype)
+
+                self.data.append((img_path, dict(boxes=box_targets, labels=labels)))
 
     def __getitem__(self, index: int) -> Tuple[np.ndarray, Dict[str, Any]]:
         img, target = self._read_sample(index)

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -131,14 +131,13 @@ def mock_ic13(tmpdir_factory, mock_image_stream):
     file = BytesIO(mock_image_stream)
     image_folder = tmpdir_factory.mktemp("images")
     label_folder = tmpdir_factory.mktemp("labels")
-    labels = ["0,168,48,166,43,219,2,216,Symbols,-\n",
-              "102,131,951,142,958,251,126,255,Latin,TENTURE\n",
-              "1030,146,1745,142,1747,241,1029,246,Latin,MURALE\n",
-              "72,653,786,648,787,793,85,812,Latin,L'atelier\n",
-              "839,648,1226,650,1228,784,843,788,Latin,Kern\n"]
-    for i in range(10):
-        # labels needs same name as image
-        fn_l = label_folder.join(f"mock_image_file_{i}.txt")
+    labels = ["1309, 2240, 1440, 2341, 'I'\n",
+              "800, 2240, 1440, 2341, 'am'\n",
+              "500, 2240, 1440, 2341, 'a'\n",
+              "900, 2240, 1440, 2341, 'jedi'\n",
+              "400, 2240, 1440, 2341, '!'"]
+    for i in range(5):
+        fn_l = label_folder.join(f"gt_mock_image_file_{i}.txt")
         with open(fn_l, 'w') as f:
             f.writelines(labels)
         fn_i = image_folder.join(f"mock_image_file_{i}.jpg")

diff --git a/tests/pytorch/test_datasets_pt.py b/tests/pytorch/test_datasets_pt.py
@@ -204,18 +204,18 @@ def test_charactergenerator():
 
 
 @pytest.mark.parametrize(
-    "train, size, rotate",
+    "size, rotate",
     [
-        [True, 9, True],  # Actual set has 9000 samples
-        [False, 1, False],  # Actual set has 1000 samples
+        [5, True],  # Actual set has 229 train and 233 test samples
+        [5, False]
+
     ],
 )
-def test_ic13_dataset(mock_ic13, train, size, rotate):
+def test_ic13_dataset(mock_ic13, size, rotate):
     input_size = (512, 512)
     ds = datasets.IC13(
         *mock_ic13,
         sample_transforms=Resize(input_size),
-        train=train,
         rotated_bbox=rotate,
     )
 

diff --git a/tests/tensorflow/test_datasets_tf.py b/tests/tensorflow/test_datasets_tf.py
@@ -190,20 +190,21 @@ def test_charactergenerator():
 
 
 @pytest.mark.parametrize(
-    "train, size, rotate",
+    "size, rotate",
     [
-        [True, 9, True],  # Actual set has 9000 samples
-        [False, 1, False],  # Actual set has 1000 samples
+        [5, True],  # Actual set has 229 train and 233 test samples
+        [5, False]
+
     ],
 )
-def test_icdar19_dataset(mock_ic13, train, size, rotate):
+def test_ic13_dataset(mock_ic13, size, rotate):
     input_size = (512, 512)
     ds = datasets.IC13(
         *mock_ic13,
         sample_transforms=Resize(input_size),
-        train=train,
         rotated_bbox=rotate,
     )
+
     assert len(ds) == size
     img, target = ds[0]
     assert isinstance(img, tf.Tensor)