Views
No views yet

pip install numpy opencv-python requests pillow transformers tensorflow1import os
2os.environ["KERAS_BACKEND"] = "tensorflow"
3
4import keras
5import numpy as np
6import cv2
7import requests
8from PIL import Image
9from io import BytesIO
10from typing import List, Optional
11from huggingface_hub import hf_hub_download
12import tensorflow as tf
13import pickle
14
15class ImageTokenizer:
16 def __init__(self):
17 self.unique_pixels = set()
18 self.pixel_to_token = {}
19 self.token_to_pixel = {}
20
21 def fit(self, images):
22 for image in images:
23 self.unique_pixels.update(np.unique(image))
24 self.pixel_to_token = {pixel: i for i, pixel in enumerate(sorted(self.unique_pixels))}
25 self.token_to_pixel = {i: pixel for pixel, i in self.pixel_to_token.items()}
26
27 def tokenize(self, images):
28 return np.vectorize(self.pixel_to_token.get)(images)
29
30 def detokenize(self, tokens):
31 return np.vectorize(self.token_to_pixel.get)(tokens)
32
33class MNISTPredictor:
34 def __init__(self, model_name):
35 # Download the model and tokenizer files
36 model_path = hf_hub_download(repo_id=model_name, filename="mnist_model.keras")
37 tokenizer_path = hf_hub_download(repo_id=model_name, filename="mnist_tokenizer.pkl")
38
39 # Load the model and tokenizer
40 self.model = keras.models.load_model(model_path)
41 with open(tokenizer_path, 'rb') as tokenizer_file:
42 self.tokenizer = pickle.load(tokenizer_file)
43
44 def extract_features(self, image: Image.Image) -> List[np.ndarray]:
45 """Extract features from the image for multiple digits."""
46 # Convert to grayscale
47 gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
48
49 # Apply Gaussian blur
50 blurred = cv2.GaussianBlur(gray, (5, 5), 0)
51
52 # Apply adaptive thresholding
53 thresh = cv2.adaptiveThreshold(blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 11, 2)
54
55 # Find contours
56 contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
57
58 digit_images = []
59 for contour in contours:
60 # Filter small contours
61 if cv2.contourArea(contour) > 50: # Adjust this threshold as needed
62 x, y, w, h = cv2.boundingRect(contour)
63 roi = thresh[y:y+h, x:x+w]
64 resized = cv2.resize(roi, (28, 28), interpolation=cv2.INTER_AREA)
65 digit_images.append(resized.reshape((28, 28, 1)).astype('float32') / 255)
66
67 return digit_images
68
69 def predict(self, image: Image.Image) -> Optional[List[int]]:
70 """Predict digits in the image."""
71 try:
72 digit_images = self.extract_features(image)
73 tokenized_images = [self.tokenizer.tokenize(img) for img in digit_images]
74 predictions = self.model.predict(np.array(tokenized_images), verbose=0)
75 return np.argmax(predictions, axis=1).tolist()
76 except Exception as e:
77 print(f"Error during prediction: {e}")
78 return None
79
80def download_image(url: str) -> Optional[Image.Image]:
81 """Download an image from a URL."""
82 try:
83 response = requests.get(url)
84 response.raise_for_status()
85 return Image.open(BytesIO(response.content))
86 except Exception as e:
87 print(f"Error downloading image: {e}")
88 return None
89
90def save_predictions_to_file(predictions: List[int], output_path: str) -> None:
91 """Save predictions to a text file."""
92 try:
93 with open(output_path, 'w') as f:
94 f.write(f"Predicted digits are: {', '.join(map(str, predictions))}\n")
95 except Exception as e:
96 print(f"Error saving predictions to file: {e}")
97
98def main(image_url: str, model_name: str, output_path: str) -> None:
99 try:
100 predictor = MNISTPredictor(model_name)
101
102 # Download image
103 image = download_image(image_url)
104 if image is None:
105 raise Exception("Failed to download image")
106
107 print(f"Image downloaded successfully.")
108
109 # Predict digits
110 digits = predictor.predict(image)
111 if digits is not None:
112 print(f"Predicted digits are: {digits}")
113
114 # Save predictions to file
115 save_predictions_to_file(digits, output_path)
116 print(f"Predictions saved to {output_path}")
117 else:
118 print("Failed to predict digits.")
119 except Exception as e:
120 print(f"An error occurred: {e}")
121
122if __name__ == "__main__":
123 image_url = "https://miro.medium.com/v2/resize:fit:720/format:webp/1*w7pBsjI3t3ZP-4Gdog-JdQ.png"
124 model_name = "0xnu/mnist-ocr"
125 output_path = "predictions.txt"
126
127 main(image_url, model_name, output_path)