update app and include examples

app.py

@@ -1,34 +1,34 @@
+import warnings
+from fastai.vision.all import *
 import gradio as gr
-import torch
-from PIL import Image
-import torchvision.transforms as transforms
 
-model = torch.load("model.pkl", map_location=torch.device("cpu"))
-model.eval()
+warnings.simplefilter("ignore", category=UserWarning)
 
+learn = load_learner("model.pkl")
 
-def preprocess_image(image):
-    transform = transforms.Compose(
-        [
-            transforms.Resize((224, 224)),
-            transforms.ToTensor(),
-            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
-        ]
-    )
-    image_tensor = transform(image)
-    return image_tensor.unsqueeze(0)
+warnings.resetwarnings()
 
+categories = ("heavy truck", "car", "bike")
 
-def classify_image(image):
-    input_tensor = preprocess_image(image)
 
-    with torch.no_grad():
-        outputs = model(input_tensor)
-        predictions = torch.argmax(outputs, dim=1).item()
+def classify_image(img):
+    _, _, probs = learn.predict(img)
+    return dict(zip(categories, map(float, probs)))
 
-    return f"Predicted class: {predictions}"
 
+image = gr.Image()
+label = gr.Label()
 
-iface = gr.Interface(fn=classify_image, inputs="image", outputs="text")
+examples = [
+    "caminhao1.jpg",
+    "caminhao2.jpg",
+    "caminhao3.jpg",
+    "carro1.jpg",
+    "carro2.jpg",
+    "carro3.jpg",
+]
 
-iface.launch()
+intf = gr.Interface(fn=classify_image, inputs=image, outputs=label, examples=examples)
+
+
+intf.launch(inline=False)

requirements.txt

@@ -1,4 +1,2 @@
 gradio==2.1.2
-torch 
-torchvision
 fastai