8 files changed, 38 insertions, 19 deletions
diff --git a/src/text_recognizer/tests/support/create_emnist_lines_support_files.py b/src/text_recognizer/tests/support/create_emnist_lines_support_files.py
index 4496e40..9abe143 100644
--- a/src/text_recognizer/tests/support/create_emnist_lines_support_files.py
+++ b/src/text_recognizer/tests/support/create_emnist_lines_support_files.py
@@ -1,4 +1,6 @@
 """Module for creating EMNIST Lines test support files."""
+# flake8: noqa: S106
+
 from pathlib import Path
 import shutil
 
@@ -17,23 +19,31 @@ def create_emnist_lines_support_files() -> None:
     SUPPORT_DIRNAME.mkdir()
 
     # TODO: maybe have to add args to dataset.
-    dataset = EmnistLinesDataset()
+    dataset = EmnistLinesDataset(
+        init_token="<sos>",
+        pad_token="_",
+        eos_token="<eos>",
+        transform=[{"type": "ToTensor", "args": {}}],
+        target_transform=[
+            {
+                "type": "AddTokens",
+                "args": {"init_token": "<sos>", "pad_token": "_", "eos_token": "<eos>"},
+            }
+        ],
+    )  # nosec: S106
     dataset.load_or_generate_data()
 
-    for index in [0, 1, 3]:
+    for index in [5, 7, 9]:
         image, target = dataset[index]
+        if len(image.shape) == 3:
+            image = image.squeeze(0)
         print(image.sum(), image.dtype)
 
-        label = (
-            "".join(
-                dataset.mapper[label]
-                for label in np.argmax(target[1:], dim=-1).flatten()
-            )
-            .stip()
-            .strip(dataset.mapper.pad_token)
+        label = "".join(dataset.mapper(label) for label in target[1:]).strip(
+            dataset.mapper.pad_token
         )
-
         print(label)
+        image = image.numpy()
         util.write_image(image, str(SUPPORT_DIRNAME / f"{label}.png"))
 
 
diff --git a/src/text_recognizer/tests/support/create_iam_lines_support_files.py b/src/text_recognizer/tests/support/create_iam_lines_support_files.py
index bb568ee..50f9e3d 100644
--- a/src/text_recognizer/tests/support/create_iam_lines_support_files.py
+++ b/src/text_recognizer/tests/support/create_iam_lines_support_files.py
@@ -1,4 +1,5 @@
 """Module for creating IAM Lines test support files."""
+# flake8: noqa
 from pathlib import Path
 import shutil
 
@@ -17,23 +18,31 @@ def create_emnist_lines_support_files() -> None:
     SUPPORT_DIRNAME.mkdir()
 
     # TODO: maybe have to add args to dataset.
-    dataset = IamLinesDataset()
+    dataset = IamLinesDataset(
+        init_token="<sos>",
+        pad_token="_",
+        eos_token="<eos>",
+        transform=[{"type": "ToTensor", "args": {}}],
+        target_transform=[
+            {
+                "type": "AddTokens",
+                "args": {"init_token": "<sos>", "pad_token": "_", "eos_token": "<eos>"},
+            }
+        ],
+    )
     dataset.load_or_generate_data()
 
     for index in [0, 1, 3]:
         image, target = dataset[index]
+        if len(image.shape) == 3:
+            image = image.squeeze(0)
         print(image.sum(), image.dtype)
 
-        label = (
-            "".join(
-                dataset.mapper[label]
-                for label in np.argmax(target[1:], dim=-1).flatten()
-            )
-            .stip()
-            .strip(dataset.mapper.pad_token)
+        label = "".join(dataset.mapper(label) for label in target[1:]).strip(
+            dataset.mapper.pad_token
         )
-
         print(label)
+        image = image.numpy()
         util.write_image(image, str(SUPPORT_DIRNAME / f"{label}.png"))
 
 
diff --git a/src/text_recognizer/tests/support/emnist_lines/Knox Ky<eos>.png b/src/text_recognizer/tests/support/emnist_lines/Knox Ky<eos>.png
new file mode 100644
index 0000000..b7d0618
--- /dev/null
+++ b/src/text_recognizer/tests/support/emnist_lines/Knox Ky<eos>.png
diff --git a/src/text_recognizer/tests/support/emnist_lines/ancillary beliefs and<eos>.png b/src/text_recognizer/tests/support/emnist_lines/ancillary beliefs and<eos>.png
new file mode 100644
index 0000000..14a8cf3
--- /dev/null
+++ b/src/text_recognizer/tests/support/emnist_lines/ancillary beliefs and<eos>.png
diff --git a/src/text_recognizer/tests/support/emnist_lines/they<eos>.png b/src/text_recognizer/tests/support/emnist_lines/they<eos>.png
new file mode 100644
index 0000000..7f05951
--- /dev/null
+++ b/src/text_recognizer/tests/support/emnist_lines/they<eos>.png
diff --git a/src/text_recognizer/tests/support/iam_lines/He rose from his breakfast-nook bench<eos>.png b/src/text_recognizer/tests/support/iam_lines/He rose from his breakfast-nook bench<eos>.png
new file mode 100644
index 0000000..6eeb642
--- /dev/null
+++ b/src/text_recognizer/tests/support/iam_lines/He rose from his breakfast-nook bench<eos>.png
diff --git a/src/text_recognizer/tests/support/iam_lines/and came into the livingroom, where<eos>.png b/src/text_recognizer/tests/support/iam_lines/and came into the livingroom, where<eos>.png
new file mode 100644
index 0000000..4974cf8
--- /dev/null
+++ b/src/text_recognizer/tests/support/iam_lines/and came into the livingroom, where<eos>.png
diff --git a/src/text_recognizer/tests/support/iam_lines/his entrance. He came, almost falling<eos>.png b/src/text_recognizer/tests/support/iam_lines/his entrance. He came, almost falling<eos>.png
new file mode 100644
index 0000000..a731245
--- /dev/null
+++ b/src/text_recognizer/tests/support/iam_lines/his entrance. He came, almost falling<eos>.png