summaryrefslogtreecommitdiff
path: root/src/text_recognizer/datasets/data_loader.py
diff options
context:
space:
mode:
authoraktersnurra <gustaf.rydholm@gmail.com>2020-06-23 22:39:54 +0200
committeraktersnurra <gustaf.rydholm@gmail.com>2020-06-23 22:39:54 +0200
commit7c4de6d88664d2ea1b084f316a11896dde3e1150 (patch)
treecbde7e64c8064e9b523dfb65cd6c487d061ec805 /src/text_recognizer/datasets/data_loader.py
parenta7a9ce59fc37317dd74c3a52caf7c4e68e570ee8 (diff)
latest
Diffstat (limited to 'src/text_recognizer/datasets/data_loader.py')
-rw-r--r--src/text_recognizer/datasets/data_loader.py15
1 files changed, 15 insertions, 0 deletions
diff --git a/src/text_recognizer/datasets/data_loader.py b/src/text_recognizer/datasets/data_loader.py
new file mode 100644
index 0000000..fd55934
--- /dev/null
+++ b/src/text_recognizer/datasets/data_loader.py
@@ -0,0 +1,15 @@
+"""Data loader collection."""
+
+from typing import Dict
+
+from torch.utils.data import DataLoader
+
+from text_recognizer.datasets.emnist_dataset import fetch_emnist_data_loader
+
+
+def fetch_data_loader(data_loader_args: Dict) -> DataLoader:
+ """Fetches the specified PyTorch data loader."""
+ if data_loader_args.pop("name").lower() == "emnist":
+ return fetch_emnist_data_loader(data_loader_args)
+ else:
+ raise NotImplementedError