Update app.py

app.py

@@ -7,6 +7,8 @@ import torchvision.transforms as transforms
 import pickle
 from torchvision import transforms, models
 
+device = torch.device('cpu')
+
 class FineTunedVGG(nn.Module):
     def __init__(self, num_classes, input_size=224):
         super(FineTunedVGG, self).__init__()
@@ -50,7 +52,7 @@ model = torch.load("model.pth",map_location ='cpu')
 with open("encoder.pkl", "rb") as encoder_file:
     label_encoder = pickle.load(encoder_file)
     
-def preprocess_image(image_path):
+def preprocess_image(image):
     transform = transforms.Compose([
         transforms.Resize((224, 224)), 
         transforms.ToTensor(),         
@@ -63,11 +65,23 @@ def preprocess_image(image_path):
 
 
 def recognize_image(image):
-    output = model.predict(image)
-    probs = torch.softmax(output, dim=1)[0].tolist()  
-    class_labels = label_encoder.classes_  
-    output_dict = dict(zip(class_labels, map(float, probs)))
-    return output_dict
+   mean = [0.0, 0.0, 0.0]
+   std = [1.0, 1.0, 1.0]
+   transform_norm = transforms.Compose([transforms.ToTensor(), 
+   transforms.Resize((224,224)),transforms.Normalize(mean, std)])
+
+   img_normalized = transform_norm(image).float()
+   img_normalized = img_normalized.unsqueeze_(0)
+
+   img_normalized = img_normalized.to(device)
+
+   with torch.no_grad():
+      model.eval()  
+      output =model(img_normalized)
+      probs = torch.softmax(output, dim=1)[0].tolist()  
+      class_labels = label_encoder.classes_  
+      output_dict = dict(zip(class_labels, map(float, probs)))
+      return output_dict
     
 image = gr.inputs.Image(shape=(224,224))
 label = gr.outputs.Label(num_top_classes=10)
@@ -81,4 +95,4 @@ examples = [
     
     
 iface = gr.Interface(fn=recognize_image, inputs=image, outputs=label, examples=examples)
-iface.launch(inline=False)
+iface.launch(inline=False)