Working on EMNIST Lines dataset

1 files changed, 30 insertions, 0 deletions

diff --git a/tests/support/create_emnist_support_files.py b/tests/support/create_emnist_support_files.py
new file mode 100644
index 0000000..f9ff030
--- /dev/null
+++ b/tests/support/create_emnist_support_files.py
@@ -0,0 +1,30 @@
+"""Module for creating EMNIST test support files."""
+from pathlib import Path
+import shutil
+
+from text_recognizer.datasets import EmnistDataset
+from text_recognizer.util import write_image
+
+SUPPORT_DIRNAME = Path(__file__).parents[0].resolve() / "emnist"
+
+
+def create_emnist_support_files() -> None:
+    """Create support images for test of CharacterPredictor class."""
+    shutil.rmtree(SUPPORT_DIRNAME, ignore_errors=True)
+    SUPPORT_DIRNAME.mkdir()
+
+    dataset = EmnistDataset(train=False)
+    dataset.load_or_generate_data()
+
+    for index in [5, 7, 9]:
+        image, label = dataset[index]
+        if len(image.shape) == 3:
+            image = image.squeeze(0)
+        image = image.numpy()
+        label = dataset.mapper(int(label))
+        print(index, label)
+        write_image(image, str(SUPPORT_DIRNAME / f"{label}.png"))
+
+
+if __name__ == "__main__":
+    create_emnist_support_files()