Update app.py

app.py

@@ -3,12 +3,13 @@ import torch
 import torch.nn as nn
 from torchvision import transforms
 from PIL import Image
+import io
 import logging
 
-# Set up logging
+# Set up logging for debugging
 logging.basicConfig(level=logging.DEBUG)
 
-# Define the model architecture
+# Define the model architecture that matches the saved .pth file
 class BacterialMorphologyClassifier(nn.Module):
     def __init__(self):
         super(BacterialMorphologyClassifier, self).__init__()
@@ -34,51 +35,48 @@ class BacterialMorphologyClassifier(nn.Module):
         x = self.fc(x)
         return x
 
-# Load the model
+# Load the model and weights
 MODEL_PATH = "https://huggingface.co/yolac/BacterialMorphologyClassification/resolve/main/model.pth"
+logging.debug("Starting model loading...")
 model = BacterialMorphologyClassifier()
-try:
-    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.info("Model loaded successfully.")
-except Exception as e:
-    logging.error(f"Error loading the model: {e}")
-    raise
+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.")
 
-# Image preprocessing transformations
+# Define image preprocessing transformations
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
 ])
 
+# Define the prediction function
 def predict(image):
     try:
-        logging.info("Received image for prediction.")
+        # Preprocess the image
         image_tensor = transform(image).unsqueeze(0)
         
         # Make prediction
         output = model(image_tensor)
         prediction = output.argmax().item()
-        confidence = output.max().item()
-
-        logging.debug(f"Model output: {output}, Prediction: {prediction}, Confidence: {confidence}")
-
+        
         # Class mapping
         class_labels = {0: 'cocci', 1: 'bacilli', 2: 'spirilla'}
-        return class_labels[prediction], confidence
-
+        
+        # Log prediction details
+        logging.debug(f"Predicted class: {class_labels[prediction]}, Confidence: {output.max().item()}")
+        
+        # Return prediction result
+        return class_labels[prediction], float(output.max().item())
     except Exception as e:
-        logging.error(f"Error during prediction: {e}")
-        return {'error': str(e)}
+        logging.error(f"Error during prediction: {str(e)}")
+        return "Error", 0.0
 
-# Create Gradio app
-gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil", label="Upload an image"),
-    outputs=["text", "number"],
-    examples=[
+# Create a Gradio interface
+inputs = gr.Image(type="pil", label="Upload an image")
+outputs = gr.Label(num_top_classes=3, label="Predicted Class")
+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"]