Update app.py

app.py

@@ -1,30 +1,22 @@
-import gradio as gr
-import torch
-from torchvision.transforms import ToTensor
-from PIL import Image
-
 # Load your PyTorch model
-model = torch.load("model.pth",map_location=torch.device('cpu'))
+model = resnet50(pretrained=False)
+model.fc = nn.Linear(model.fc.in_features, 2)
+model.load_state_dict(torch.load("model.pth", map_location=torch.device('cpu')))
 
 classes = ['bom', 'ruim']
 
-def preprocess(image):
-    image = image.resize((224, 224))
-    image_tensor = ToTensor()(image)
-    image_tensor = image_tensor.unsqueeze(0)
-    return image_tensor
-
 # Define the function for image classification
 def classify_image(image):
-    image_tensor = preprocess(image)
+    image_tensor = ToTensor()(image).unsqueeze(0)
 
     # Perform inference using your PyTorch model
     with torch.no_grad():
         model.eval()
         outputs = model(image_tensor)
 
-    predicted_labels = outputs.argmax(dim=1).tolist()
-    return predicted_labels
+        _, predicted = torch.max(outputs.data, 1)
+        return classes[predicted.item()]
+
 
 # Define the Gradio interface
 inputs = gr.Image()
@@ -32,5 +24,4 @@ outputs = gr.Label(num_top_classes=1)
 
 interface = gr.Interface(fn=classify_image, inputs=inputs, outputs=outputs)
 
-# Launch the interface
 interface.launch(debug=True)