Update app.py

app.py

@@ -3,6 +3,7 @@ import torch
 import torch.nn as nn
 from torchvision import transforms
 from PIL import Image
+import io
 import logging
 
 # Set up logging for debugging
@@ -37,14 +38,11 @@ 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...")
-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)}")
+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.")
 
 # Define image preprocessing transformations
 transform = transforms.Compose([
@@ -56,10 +54,8 @@ 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)
@@ -78,14 +74,8 @@ def predict(image):
         return "Error", 0.0
 
 # Create a Gradio interface
-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(debug=True)
+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)