update app.py

app.py

@@ -1,7 +1,44 @@
-import gradio as gr
+import gradio as gr 
+import model_builder as mb 
+from torchvision import transforms
+import torch 
 
-def greet(name):
-    return "Hello " + name + "!!"
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
+manual_transform = transforms.Compose([
+    transforms.ToPILImage(), 
+    transforms.Resize(size=(224, 224)), 
+    transforms.ToTensor(), 
+    normalize
+])
+
+class_names = ['Fresh Banana',
+  'Fresh Lemon',
+  'Fresh Lulo',
+  'Fresh Mango',
+  'Fresh Orange',
+  'Fresh Strawberry',
+  'Fresh Tamarillo',
+  'Fresh Tomato',
+  'Spoiled Banana',
+  'Spoiled Lemon',
+  'Spoiled Lulo',
+  'Spoiled Mango',
+  'Spoiled Orange',
+  'Spoiled Strawberry',
+  'Spoiled Tamarillo',
+  'Spoiled Tomato']
+
+model_0 = mb.create_model_baseline_effnetb0(out_feats=len(class_names), device=device)
+model_0.load_state_dict(torch.load(f="models/effnetb0_fruitsvegs0_5_epochs.pt", weights_only=True))
+
+def pred(img):
+    model_0.eval()
+    transformed = manual_transform(img).to(device)
+    with torch.inference_mode():
+        logits = model_0(transformed.unsqueeze(dim=0))
+        pred = torch.softmax(logits, dim=-1)
+    return f"prediction: {class_names[pred.argmax(dim=-1).item()]} | confidence: {pred.max():.3f}"
+
+demo = gr.Interface(pred, gr.Image(), outputs="text")
 demo.launch()