fix

app.py

@@ -1,21 +1,39 @@
 import gradio as gr
+import torch
+import torchvision.transforms as transforms
+import torch.nn as nn
+import numpy as np
 
 def greet(name):
     return "Hello " + name + "!!"
 
 def predict(model, image, device):
+	# Preprocess the image
+    transform = transforms.Compose([
+        transforms.Resize(256),
+        transforms.CenterCrop(224),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])
+    img_tensor = transform(image).unsqueeze(0)
+
     pred_fun = torch.nn.Softmax(dim=1)
-    preds = []
+    # preds = []
     with torch.set_grad_enabled(False):
-        x = image.to(device)
-        y = pred_fun(model(x))
+        y = pred_fun(model(img_tensor))
+        print(y)
     y = y.cpu().numpy()
+    print(y)
     y = y[:, 1]  # cat:0, dog: 1
-    preds.append(y)
-    preds = np.concatenate(preds)
-    return preds
+    print(y)
+    y = y[0]
+    print(y)
+    # preds.append(y)
+    # preds = np.concatenate(preds)
+    return {"tenka ippin":y, "no entry":1-y}
+    # return preds
 
-def process_image(input_image, output_type):
+def process_image(input_image):
 	model = torch.load('models/tenichi_noentry.pth')
 	preds = predict(model, input_image, 'cpu')
 	return preds