Update app.py

app.py

@@ -4,8 +4,12 @@ import torch.nn as nn
 from torchvision import transforms
 from PIL import Image
 import io
+import logging
 
-# Define the model architecture that matches the saved .pth file
+# Set up logging
+logging.basicConfig(level=logging.DEBUG)
+
+# Define the model architecture
 class BacterialMorphologyClassifier(nn.Module):
     def __init__(self):
         super(BacterialMorphologyClassifier, self).__init__()
@@ -31,56 +35,53 @@ class BacterialMorphologyClassifier(nn.Module):
         x = self.fc(x)
         return x
 
-# Load the model and weights
-model = BacterialMorphologyClassifier()
+# Load the model
 MODEL_PATH = "https://huggingface.co/yolac/BacterialMorphologyClassification/resolve/main/model.pth"
-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()
+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
 
-# Define image preprocessing transformations
+# 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 Gradio interface
 def predict(image):
     try:
-        # Preprocess the image
+        logging.info("Received image for prediction.")
         image_tensor = transform(image).unsqueeze(0)
         
         # Make prediction
         output = model(image_tensor)
         prediction = output.argmax().item()
         confidence = output.max().item()
-        
-        # Print debugging information
-        print(f"Predicted: {prediction}, Confidence: {confidence}")
-        print(f"Model Output: {output}")
+
+        logging.debug(f"Model output: {output}, Prediction: {prediction}, Confidence: {confidence}")
 
         # Class mapping
         class_labels = {0: 'cocci', 1: 'bacilli', 2: 'spirilla'}
-        
-        # Return prediction result
         return class_labels[prediction], confidence
+
     except Exception as e:
+        logging.error(f"Error during prediction: {e}")
         return {'error': str(e)}
 
-
-# Create the Gradio interface
-iface = gr.Interface(
+# Create Gradio app
+gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="pil", label="Upload an image"),
-    outputs=[gr.Label(num_top_classes=3, label="Predicted Class"), gr.Number(label="Confidence")],
+    inputs=gr.inputs.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"
+        ["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 the app
-if __name__ == "__main__":
-    iface.launch()
+).launch(debug=True)