Upload 4 files

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import torch
+import numpy as np
 from PIL import Image
 from torchvision.transforms import Compose, Resize, ToTensor, Normalize
 
@@ -19,17 +20,34 @@ transform = Compose([
 ])
 
 def predict(img):
-    img = Image.fromarray(img.astype('uint8'), 'RGB')
+    labels = list(range(10))
+    if isinstance(img, np.ndarray):
+        img = Image.fromarray(img.astype('uint8'), 'RGB')
     img = transform(img)
     img = img.unsqueeze(0)
 
-    prediction = model(img).argmax(axis=1)
-    return f'Model prediction is {prediction[0]}'
+    with torch.inference_mode():
+        prediction = torch.softmax(model(img),dim=1)[0]
+
+
+    result = { num:float(prob.numpy()) for num, prob in enumerate(prediction)}
+
+    
+    return result
+
+
+example_images = [
+    "1.png",  # Make sure these paths are correct
+    "2.png",
+    "3.png"
+]
+
 
 demo = gr.Interface(
     fn=predict,
     inputs=["image"],
-    outputs=["text"],
+    outputs=[gr.Label(num_top_classes=5, label="Predictions")],
+    examples=example_images,    
 )
 
 demo.launch()