Create app.py

app.py

@@ -0,0 +1,53 @@
+import numpy as np
+import gradio as gr
+from PIL import Image
+import torch
+import torch.nn.functional as F
+
+# Load the TorchScript model (make sure to place 'lenet_cnn.pth' in your repository or accessible path)
+model = torch.jit.load("lenet_cnn.pth")
+model.eval()
+
+# Set the device (here we assume CPU, adjust if needed)
+device = torch.device("cpu")
+
+def predict(data):
+    try:
+        # Extract the drawn image from the input
+        image = data["composite"]
+        if image is None or np.sum(image) == 0:
+            return "Error: No strokes detected. Please draw a digit."
+        
+        # Convert to grayscale using the alpha channel and resize to 28x28
+        image = Image.fromarray(image[:, :, 3])
+        image = image.resize((28, 28)).convert("L")
+        
+        # Normalize and convert to tensor
+        image = np.array(image, dtype=np.float32) / 255.0
+        image = torch.tensor(image, dtype=torch.float32).unsqueeze(0).unsqueeze(0)
+        image = image.to(device)
+        
+        # Run the model inference
+        with torch.no_grad():
+            output = model(image)
+            probabilities = F.softmax(output, dim=1).squeeze(0).tolist()
+        
+        # Create a dictionary mapping digit to probability (as a percentage)
+        result = {str(i): prob * 100 for i, prob in enumerate(probabilities)}
+        return result
+
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+# Create the Gradio Interface
+interface = gr.Interface(
+    fn=predict,
+    inputs=gr.Sketchpad(width=560, height=560, brush=gr.Brush(line_width=25)),  # Using "line_width" to adjust brush size
+    outputs=gr.Label(num_top_classes=3),
+    title="LeNet Handwritten Digit Classifier",
+    description="Draw a digit and press 'Submit' to classify it.",
+    theme="dark"
+)
+
+if __name__ == "__main__":
+    interface.launch(share=True)