Update app.py

app.py

@@ -3,7 +3,7 @@ import torchvision
 import torch.nn as nn
 import torchvision.transforms as transforms
 
-model = torchvision.models.resnet50(pretrained=False)
+model = torchvision.models.resnet50(pretrained=True)
 model.fc = nn.Linear(model.fc.in_features, 2)
 model.load_state_dict(torch.load("model.pth", map_location=torch.device('cpu')))
 model.eval()
@@ -15,6 +15,8 @@ transform = transforms.Compose([
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
 ])
 
+classes = ['Fruta própria para o consumo', 'Fruta imprópria para o consumo']
+
 import gradio as gr
 from PIL import Image
 
@@ -25,7 +27,7 @@ def predict(image):
     with torch.no_grad():
         output = model(image)
         _, predicted = torch.max(output.data, 1)
-        return dataset.classes[predicted.item()]
+        return classes[predicted.item()]
 
 # Define the input and output components
 image_input = gr.inputs.Image(type="pil", label="Upload Image")