Update README.md

README.md

@@ -20,19 +20,50 @@ pip install numpy opencv-python requests pillow transformers tensorflow
 ### Usage
 
 ```python
+import os
+os.environ["KERAS_BACKEND"] = "tensorflow"
+
+import keras
 import numpy as np
 import cv2
 import requests
 from PIL import Image
 from io import BytesIO
-from transformers import TFAutoModelForImageClassification, AutoFeatureExtractor
+from typing import List, Optional
+from huggingface_hub import hf_hub_download
+import tensorflow as tf
+import pickle
+
+class ImageTokenizer:
+    def __init__(self):
+        self.unique_pixels = set()
+        self.pixel_to_token = {}
+        self.token_to_pixel = {}
+
+    def fit(self, images):
+        for image in images:
+            self.unique_pixels.update(np.unique(image))
+        self.pixel_to_token = {pixel: i for i, pixel in enumerate(sorted(self.unique_pixels))}
+        self.token_to_pixel = {i: pixel for pixel, i in self.pixel_to_token.items()}
+
+    def tokenize(self, images):
+        return np.vectorize(self.pixel_to_token.get)(images)
+
+    def detokenize(self, tokens):
+        return np.vectorize(self.token_to_pixel.get)(tokens)
 
 class MNISTPredictor:
     def __init__(self, model_name):
-        self.model = TFAutoModelForImageClassification.from_pretrained(model_name)
-        self.feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
+        # Download the model and tokenizer files
+        model_path = hf_hub_download(repo_id=model_name, filename="mnist_model.keras")
+        tokenizer_path = hf_hub_download(repo_id=model_name, filename="mnist_tokenizer.pkl")
+
+        # Load the model and tokenizer
+        self.model = keras.models.load_model(model_path)
+        with open(tokenizer_path, 'rb') as tokenizer_file:
+            self.tokenizer = pickle.load(tokenizer_file)
 
-    def extract_features(self, image):
+    def extract_features(self, image: Image.Image) -> List[np.ndarray]:
         """Extract features from the image for multiple digits."""
         # Convert to grayscale
         gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
@@ -53,26 +84,22 @@ class MNISTPredictor:
                 x, y, w, h = cv2.boundingRect(contour)
                 roi = thresh[y:y+h, x:x+w]
                 resized = cv2.resize(roi, (28, 28), interpolation=cv2.INTER_AREA)
-                digit_images.append(Image.fromarray(resized).convert('RGB'))
+                digit_images.append(resized.reshape((28, 28, 1)).astype('float32') / 255)
 
         return digit_images
 
-    def predict(self, image):
+    def predict(self, image: Image.Image) -> Optional[List[int]]:
         """Predict digits in the image."""
         try:
             digit_images = self.extract_features(image)
-            predictions = []
-            for digit_image in digit_images:
-                inputs = self.feature_extractor(images=digit_image, return_tensors="tf")
-                outputs = self.model(**inputs)
-                predicted_class = int(np.argmax(outputs.logits))
-                predictions.append(predicted_class)
-            return predictions
+            tokenized_images = [self.tokenizer.tokenize(img) for img in digit_images]
+            predictions = self.model.predict(np.array(tokenized_images), verbose=0)
+            return np.argmax(predictions, axis=1).tolist()
         except Exception as e:
             print(f"Error during prediction: {e}")
             return None
 
-def download_image(url):
+def download_image(url: str) -> Optional[Image.Image]:
     """Download an image from a URL."""
     try:
         response = requests.get(url)
@@ -82,7 +109,7 @@ def download_image(url):
         print(f"Error downloading image: {e}")
         return None
 
-def save_predictions_to_file(predictions, output_path):
+def save_predictions_to_file(predictions: List[int], output_path: str) -> None:
     """Save predictions to a text file."""
     try:
         with open(output_path, 'w') as f:
@@ -90,7 +117,7 @@ def save_predictions_to_file(predictions, output_path):
     except Exception as e:
         print(f"Error saving predictions to file: {e}")
 
-def main(image_url, model_name, output_path):
+def main(image_url: str, model_name: str, output_path: str) -> None:
     try:
         predictor = MNISTPredictor(model_name)
 
@@ -103,11 +130,14 @@ def main(image_url, model_name, output_path):
 
         # Predict digits
         digits = predictor.predict(image)
-        print(f"Predicted digits are: {digits}")
-
-        # Save predictions to file
-        save_predictions_to_file(digits, output_path)
-        print(f"Predictions saved to {output_path}")
+        if digits is not None:
+            print(f"Predicted digits are: {digits}")
+
+            # Save predictions to file
+            save_predictions_to_file(digits, output_path)
+            print(f"Predictions saved to {output_path}")
+        else:
+            print("Failed to predict digits.")
     except Exception as e:
         print(f"An error occurred: {e}")