Update app.py

app.py

@@ -1,25 +1,44 @@
 import gradio as gr
 import torch
-from FlowerClassificationModel import FlowerClassificationModel
+from FlowerClassificationModel import FlowerClassificationModel  # Replace with your model's class name
+from torchvision import transforms
+from PIL import Image
 
 # Load the model
-model = FlowerClassificationModel()
-model.load_state_dict(torch.load("flower_classification_model.pth"))
-model.eval()
+model = FlowerClassificationModel()  # Instantiate your model
+model.load_state_dict(torch.load("flower_classification_model.pth", map_location=torch.device('cpu')))
+model.eval()  # Set the model to evaluation mode
 
+# Define image preprocessing
+preprocess = transforms.Compose([
+    transforms.Resize((224, 224)),  # Adjust to your model's input size
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
+# Define the prediction function
 def classify_flower(image):
-    # Add your image preprocessing logic here
-    # Perform prediction with the model
-    prediction = model(image)  # Assuming this is how your model works
-    return prediction
+    # Preprocess the input image
+    image = Image.fromarray(image)  # Convert NumPy array to PIL Image
+    input_tensor = preprocess(image).unsqueeze(0)  # Add batch dimension
+
+    # Perform prediction
+    with torch.no_grad():
+        output = model(input_tensor)
+        _, predicted = torch.max(output, 1)
+    
+    # Map prediction to class label
+    labels = ["Class1", "Class2", "Class3", "Class4", "Class5"]  # Replace with your actual class names
+    return labels[predicted.item()]
 
-# Gradio Interface
-interface = gr.Interface(
+# Create the Gradio interface
+demo = gr.Interface(
     fn=classify_flower,
     inputs="image",
-    outputs="label",
+    outputs="text",
     title="Flower Classification",
-    description="Upload an image to classify the flower."
+    description="Upload an image to classify the flower type."
 )
 
-interface.launch()
+# Launch the app
+demo.launch()