Fix dataset class

3 files changed, 45 insertions, 36 deletions

diff --git a/dataset.py b/dataset.py
new file mode 100644
index 0000000..78d4691
--- /dev/null
+++ b/dataset.py
@@ -0,0 +1,45 @@
+import numpy as np
+from PIL import Image
+from pathlib import Path
+
+"""Class to interface the training and testing data."""
+
+
+class Dataset:
+    def __init__(self) -> None:
+        self.data_path = Path('./data')
+
+    """Convert the dataset to a 2 dimension array."""
+    def data(self):
+        for dir in self.data_path.iterdir():
+            if not dir.is_dir():
+                continue
+
+            for file in dir.glob('*.png'):
+                image = Image.open(str(file))
+                image_array = self._img_to_array(image)
+                # Return the image's pixel values as an array alongside
+                # the character that it represents.
+                yield (dir.name, image_array)
+
+    """
+    Get an image from the dataset.
+    """
+    def get_image(self, path: str):
+        image = Image.open(f"{self.data_path}/{path}")
+
+        return self._img_to_array(image)
+
+    def get_random_sample(self):
+        pass
+
+    """
+    Grab the image in RGB, add a white background, and return it as
+    a black and white array.
+    """
+    def _img_to_array(self, image: Image):
+        fill_color = (255, 255, 255)  # White background.
+        background = Image.new(image.mode[:-1], image.size, fill_color)
+        background.paste(image, image.split()[-1])
+
+        return np.asarray(background.convert(mode='1'))
diff --git a/dataset_utils.py b/dataset_utils.py
deleted file mode 100644
index 927a78a..0000000
--- a/dataset_utils.py
+++ /dev/null
@@ -1,32 +0,0 @@
-import numpy as np
-from PIL import Image
-from pathlib import Path
-
-"""Class to interface the training and testing data."""
-class DatasetUtils:
-    def __init__(self) -> None:
-        self.data_path = Path('./data')
-
-    """Convert the dataset to a 2 dimension array."""
-    def data_to_arrays(self):
-        for dir in self.data_path.iterdir():
-            if not dir.is_dir(): continue
-
-            for file in dir.glob('*.png'):
-                image = Image.open(str(file)).convert(mode='L')
-                # Return the image's pixel values as an array alongside
-                # the character that it represents.
-                yield (dir.name, np.asarray(image))
-
-    """
-    Get an image from the dataset.
-    """
-    def get_image(self, path: str):
-        return np.asarray(Image.open(f"{self.data_path}/{path}").convert('1'))
-                
-    """Search for a file in the dataset."""
-    def search(self, filename: str):
-        pass
-
-    def get_random_sample(self):
-        pass
diff --git a/utils.py b/utils.py
deleted file mode 100644
index a4f6cd3..0000000
--- a/utils.py
+++ /dev/null
@@ -1,4 +0,0 @@
-import numpy as np
-
-def random_array(length: int):
-    return np.random.rand(length, 1)