summaryrefslogtreecommitdiff
path: root/src/text_recognizer/datasets/data_loader.py
diff options
context:
space:
mode:
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