Update app.py

app.py

@@ -3,7 +3,6 @@ import torch
 import torch.nn as nn
 from torchvision import transforms
 from PIL import Image
-import io
 import logging
 
 # Set up logging for debugging
@@ -38,11 +37,14 @@ class BacterialMorphologyClassifier(nn.Module):
 # Load the model and weights
 MODEL_PATH = "https://huggingface.co/yolac/BacterialMorphologyClassification/resolve/main/model.pth"
 logging.debug("Starting model loading...")
-model = BacterialMorphologyClassifier()
-state_dict = torch.hub.load_state_dict_from_url(MODEL_PATH, map_location=torch.device('cpu'))
-model.load_state_dict(state_dict, strict=False)
-model.eval()
-logging.debug("Model loaded successfully.")
+try:
+    model = BacterialMorphologyClassifier()
+    state_dict = torch.hub.load_state_dict_from_url(MODEL_PATH, map_location=torch.device('cpu'))
+    model.load_state_dict(state_dict, strict=False)
+    model.eval()
+    logging.debug("Model loaded successfully.")
+except Exception as e:
+    logging.error(f"Error loading the model: {str(e)}")
 
 # Define image preprocessing transformations
 transform = transforms.Compose([
@@ -54,8 +56,10 @@ transform = transforms.Compose([
 # Define the prediction function
 def predict(image):
     try:
+        logging.debug("Starting prediction...")
         # Preprocess the image
         image_tensor = transform(image).unsqueeze(0)
+        logging.debug("Image preprocessing completed.")
         
         # Make prediction
         output = model(image_tensor)
@@ -74,8 +78,13 @@ def predict(image):
         return "Error", 0.0
 
 # Create a Gradio interface
-inputs = gr.Image(type="pil", label="Upload an image")
-outputs = gr.Label(num_top_classes=3, label="Predicted Class")
-
-# Launch the Gradio app
-gr.Interface(fn=predict, inputs=inputs, outputs=outputs, live=True).launch(server_name="0.0.0.0", server_port=7861, debug=True)
+gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil", label="Upload an image"),
+    outputs=["text", "number"],
+    examples=[
+        ["https://huggingface.co/datasets/yolac/BacterialMorphologyClassification/resolve/main/img%20290.jpg"],
+        ["https://huggingface.co/datasets/yolac/BacterialMorphologyClassification/resolve/main/img%20565.jpg"],
+        ["https://huggingface.co/datasets/yolac/BacterialMorphologyClassification/resolve/main/img%208.jpg"]
+    ]
+).launch(server_name="0.0.0.0", server_port=7862, debug=True)