Fix neural network

1 files changed, 8 insertions, 2 deletions

diff --git a/dataset_utils.py b/dataset_utils.py
index fda75d2..927a78a 100644
--- a/dataset_utils.py
+++ b/dataset_utils.py
@@ -1,4 +1,4 @@
-import numpy
+import numpy as np
 from PIL import Image
 from pathlib import Path
 
@@ -16,8 +16,14 @@ class DatasetUtils:
                 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, numpy.asarray(image).ravel())
+                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