improved

app.py

@@ -1,8 +1,12 @@
 import gradio as gr
+from gradio import Brush
 import torch
 import numpy as np
-from torchvision import transforms
+import cv2  # For Gaussian blur
 from PIL import Image
+from torch import nn, save, load
+from torchvision.transforms import Compose, ToTensor, Normalize, Resize
+
 
 # Load the model using PyTorch
 model_path = "https://huggingface.co/immartian/improved_digits_recognition/resolve/main/pytorch_model.bin"
@@ -29,44 +33,72 @@ class ImageClassifier(torch.nn.Module):
 # Instantiate the model and load weights
 model = ImageClassifier()
 model.load_state_dict(torch.hub.load_state_dict_from_url(model_path))
-model.eval()
+# model.eval()
 
-# Function to process sketchpad input
-def sketchToNumpy(image):
-    # Extract the 'composite' key from the sketchpad input dictionary
-    imArray = image['composite']  # 'composite' contains the drawn image
-    return imArray
+# with open('pytorch_model.bin', 'rb') as f: 
+#     model.load_state_dict(load(f))  
 
 # Gradio preprocessing and prediction pipeline
 def predict_digit(image):
-    # Convert the sketchpad input into a PIL Image
-    image = Image.fromarray(image).convert('L')  # Convert to grayscale
-
-    # Preprocess: resize to 28x28 and normalize
-    transform = transforms.Compose([
-        transforms.Resize((28, 28)),
-        transforms.ToTensor(),
-        transforms.Normalize((0.5,), (0.5,))
+    # Extract the 'composite' key, which contains the drawn image
+    if isinstance(image, dict):
+        image = image.get('composite', None)  # Use the composite image
+
+    if image is None:
+        raise ValueError("No image data found in the input!")
+
+    #print("Unique pixel values in the image array:", np.unique(image))
+
+    # Ensure the input is a numpy array
+    image = np.array(image, dtype=np.uint8)
+
+
+    # Apply Gaussian blur to reduce noise
+    image = cv2.GaussianBlur(image, (5, 5), 0)
+
+    # If the image has multiple channels (e.g., BGR), convert it to grayscale
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    image = cv2.resize(image, (28, 28))
+    # Optional: Apply adaptive histogram equalization to improve contrast
+    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
+    image = clahe.apply(image)
+
+
+
+    # Convert the numpy array back to a PIL Image for torchvision compatibility
+    img_pil = Image.fromarray(image)
+    img_pil.show()
+
+    transform = Compose([
+        Resize((28, 28)),
+        ToTensor(),
+        Normalize((0.5,), (0.5,))
     ])
-    
-    img_tensor = transform(image).unsqueeze(0)  # Add batch dimension
+
+    img_tensor = transform(img_pil).unsqueeze(0)  # Add batch dimension
+
+    # Debugging: Print tensor shape and some pixel values
+    print(f"Input Tensor Shape: {img_tensor.shape}")
+    print(f"First 5 pixels of the tensor: {img_tensor[0, 0, :5, :5]}")
 
     # Pass through the model
     with torch.no_grad():
         output = model(img_tensor)
-        predicted_label = torch.argmax(output, dim=1).item()
+        prediction = torch.argmax(output)
     
-    return f"Predicted Label: {predicted_label}"
+    return f"Predicted Label: {prediction}"
 
-# Create Gradio Interface
-interface = gr.Interface(
-    fn=lambda x: predict_digit(sketchToNumpy(x)),
-    inputs=gr.Sketchpad(crop_size=(256,256), type='numpy', image_mode='L'),
-    outputs="text",
-    title="Digit Recognizer",
-    description="Draw a digit (0-9) and the model will predict the number!"
-)
+# Create Gradio Interface using ImageEditor
+with gr.Blocks() as demo:
+    with gr.Row():
+        im = gr.ImageEditor(type="numpy", crop_size="1:1")
+        # im_preview = gr.Image()
+        prediction_box = gr.Textbox(label="Predicted Digit")
 
+    im.change(predict_digit, outputs=prediction_box, inputs=im, show_progress="hidden")
+
+    #im.change(predict_digit, outputs="text", inputs=im, show_progress=True)
+    
 # Launch the app
 if __name__ == "__main__":
-    interface.launch()
+    demo.launch()