Update app.py

app.py

@@ -1,9 +1,10 @@
 import gradio as gr
 import torch
-from PIL import Image
+from torchvision import transforms
+import numpy as np
 
 # Load the model
-model = torch.hub.load('pytorch/vision:v0.9.0', 'resnet101', pretrained=False)
+model = torch.hub.load('pytorch/vision:v0.9.0', 'resnet101', pretrained=True)
 model.load_state_dict(torch.load('resnet101_pneumonia.pt', map_location=torch.device('cpu')))
 model.eval()
 
@@ -14,8 +15,13 @@ class_names = ["normal", "pneumonia"]
 def predict(img):
     # Preprocess the image
     img = img.convert("RGB")
-    transform = torch.hub.load('pytorch/vision:v0.9.0', 'transforms', pretrained=True)
-    img = transform(img)
+    preprocess = 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 = preprocess(img)
     img = img.unsqueeze(0)
 
     # Make prediction
@@ -35,7 +41,7 @@ def predict(img):
 # Create the Gradio interface
 iface = gr.Interface(
     fn=predict,
-    inputs=gr.inputs.Image(type="pil", label="Input Image"),
+    inputs=gr.inputs.Image(type="numpy", label="Input Image"),
     outputs=gr.outputs.Label(num_top_classes=2, label="Predicted Class"),
     title="PneumoniaDetector 👁",
     description="A ResNet101 computer vision model to detect pneumonia",