remolve transformer

app.py

@@ -2,44 +2,62 @@ import gradio as gr
 import torch
 from torchvision import transforms
 from PIL import Image
-from transformers import AutoModelForImageClassification, AutoFeatureExtractor
 
-# Load the model and feature extractor from Hugging Face
-model_name = "immartian/improved_digits_recognition"
-model = AutoModelForImageClassification.from_pretrained(model_name)
-feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
+# Load the model using PyTorch
+model_path = "https://huggingface.co/immartian/improved_digits_recognition/resolve/main/pytorch_model.bin"
 
-# Preprocessing function to transform the drawn image into a format the model can recognize
-def preprocess_image(image):
-    # Convert the image into a format suitable for the model
-    image = Image.fromarray(image).convert('L')  # Convert to grayscale
-    image = image.resize((28, 28))  # Resize to 28x28 pixels
-    image = image.convert('RGB')  # Model expects 3-channel images, so convert to RGB
-    inputs = feature_extractor(images=image, return_tensors="pt")
-    return inputs['pixel_values']
+# Define your ImageClassifier model architecture (same as used during training)
+class ImageClassifier(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.model = torch.nn.Sequential(
+            torch.nn.Conv2d(1, 32, (3, 3)),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(32, 64, (3, 3)),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(64, 64, (3, 3)),
+            torch.nn.ReLU(),
+            torch.nn.AdaptiveAvgPool2d((1, 1)),
+            torch.nn.Flatten(),
+            torch.nn.Linear(64, 10)
+        )
+    
+    def forward(self, x):
+        return self.model(x)
+
+# Instantiate the model and load weights
+model = ImageClassifier()
+model.load_state_dict(torch.hub.load_state_dict_from_url(model_path))
+model.eval()
 
-# Prediction function to classify the drawn digit
+# Gradio preprocessing and prediction pipeline
 def predict_digit(image):
-    # Preprocess the input image
-    inputs = preprocess_image(image)
+    # Preprocess the image: resize to 28x28, convert to grayscale, and normalize
+    image = image.convert('L')  # Convert to grayscale
+    transform = transforms.Compose([
+        transforms.Resize((28, 28)),
+        transforms.ToTensor(),
+        transforms.Normalize((0.5,), (0.5,))
+    ])
     
-    # Make the prediction
-    with torch.no_grad():
-        outputs = model(inputs)
-        predicted_label = outputs.logits.argmax(-1).item()
+    img_tensor = transform(image).unsqueeze(0)  # Add batch dimension
 
-    return f"Predicted Digit: {predicted_label}"
+    # Pass through the model
+    with torch.no_grad():
+        output = model(img_tensor)
+        predicted_label = torch.argmax(output, dim=1).item()
+    
+    return f"Predicted Label: {predicted_label}"
 
-# Gradio interface for drawing the digit and displaying the prediction
-demo = gr.Interface(
+# Create Gradio Interface
+interface = gr.Interface(
     fn=predict_digit,
-    inputs="sketchpad",  # Allow users to draw a digit
+    inputs=gr.Image(source="canvas", tool="editor", type="pil"),  # User can draw on a canvas
     outputs="text",
-    title="MNIST Digit Recognition",
-    description="Draw a digit (0-9) and let the model recognize it!",
-    live=True  # The prediction updates while the user draws
+    title="Digit Recognizer",
+    description="Draw a digit (0-9) and the model will predict the number!"
 )
 
 # Launch the app
 if __name__ == "__main__":
-    demo.launch()
+    interface.launch()

requirements.txt

@@ -1,5 +1,4 @@
 torch
 torchvision
 gradio
-transformers
-Pillow  # Required for image processing
+Pillow