Update app.py

app.py

@@ -10,7 +10,7 @@ num_features = model.fc.in_features
 model.fc = torch.nn.Linear(num_features, 5)  # Replace the final layer for 5 classes
 
 # Load the model weights
-checkpoint = torch.load('shiva_flower_classification.pth', map_location=torch.device('cpu'))
+checkpoint = torch.load('shiva_flower_classification.pth', map_location=torch.device('cpu'), weights_only=True)
 
 # Get model state_dict without the 'fc' layer
 state_dict = checkpoint
@@ -37,7 +37,6 @@ transform = transforms.Compose([
 # Prediction function
 def predict(image):
     # Preprocess the image
-    image = Image.open(image).convert('RGB')
     image = transform(image).unsqueeze(0)
     
     # Predict the class
@@ -51,7 +50,7 @@ def predict(image):
 # Gradio Interface
 interface = gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="file"),
+    inputs=gr.Image(type="pil"),
     outputs="text",
     title="Flower Classification",
     description="Upload an image of a flower to classify it into one of the five categories: daisy, dandelion, rose, sunflower, or tulip."