Update app.py

app.py

@@ -1,42 +1,78 @@
+import torch
+import torch.nn as nn
 import gradio as gr
-from tensorflow.keras.models import load_model
-from huggingface_hub import hf_hub_download
-import numpy as np
+from torchvision import transforms
 from PIL import Image
+import requests
+from io import BytesIO
 
-# Define constants
-MODEL_REPO = "yolac/BacterialMorphologyClassification"
-MODEL_FILENAME = "model.keras"
-MODEL_PATH = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILENAME)
+# Define the PyTorch model architecture
+class MyPyTorchModel(nn.Module):
+    def __init__(self):
+        super(MyPyTorchModel, self).__init__()
+        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)
+        self.relu = nn.ReLU()
+        self.flatten = nn.Flatten()
+        self.fc1 = nn.Linear(32 * 224 * 224, 3)  # Adjust output size for 3 classes
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.relu(x)
+        x = self.flatten(x)
+        x = self.fc1(x)
+        return x
 
 # Load the model
-print("Loading model...")
-model = load_model(MODEL_PATH)
+model_path = 'https://huggingface.co/yolac/BacterialMorphologyClassification/resolve/main/model.pth'
+model = MyPyTorchModel()
+model.load_state_dict(torch.load(model_path))
+model.eval()
 
-# Preprocessing function
-def preprocess_image(image):
-    image = image.resize((224, 224))  # Adjust size as per your model input
-    image_array = np.array(image) / 255.0  # Normalize to [0, 1]
-    image_array = np.expand_dims(image_array, axis=0)  # Add batch dimension
-    return image_array
+# Define image transformations
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+])
 
-# Prediction function
+# Define a function to predict the class of the image
 def predict(image):
-    image_array = preprocess_image(image)
-    predictions = model.predict(image_array)
-    class_names = ["Cocci", "Bacilli", "Spirilla"]
-    predicted_class = class_names[np.argmax(predictions)]
-    return f"Predicted Class: {predicted_class}"
-
-# Gradio Interface
-interface = gr.Interface(
+    image = Image.fromarray(image).convert('RGB')
+    image = transform(image).unsqueeze(0)
+
+    with torch.no_grad():
+        outputs = model(image)
+        _, predicted = torch.max(outputs, 1)
+
+    class_labels = ['cocci', 'bacilli', 'spirilla']
+    predicted_label = class_labels[predicted.item()]
+
+    return predicted_label
+
+# URLs for 3 example images
+example_image_urls = [
+    "https://huggingface.co/datasets/yolac/BacterialMorphologyClassification/resolve/main/img%20290.jpg",  # Replace with the actual URL
+    "https://huggingface.co/datasets/yolac/BacterialMorphologyClassification/resolve/main/img%20565.jpg",  # Replace with the actual URL
+    "https://huggingface.co/datasets/yolac/BacterialMorphologyClassification/resolve/main/img%208.jpg",  # Replace with the actual URL
+]
+
+# Function to download and open example images
+def load_example_image(url):
+    response = requests.get(url)
+    image = Image.open(BytesIO(response.content))
+    return image
+
+# Load the example images
+example_images = [load_example_image(url) for url in example_image_urls]
+
+# Create a Gradio interface
+iface = gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="pil"),
+    inputs=gr.inputs.Image(shape=(224, 224), label="Upload an image or use examples"),
     outputs="text",
-    title="Bacterial Morphology Classification",
-    description="Upload an image of bacteria to classify as Cocci, Bacilli, or Spirilla."
+    title="Bacterial Morphology Classifier",
+    description="Classify images of bacteria into cocci, bacilli, or spirilla.",
+    examples=example_images  # Add example images to the Gradio app
 )
 
 # Launch the app
-if __name__ == "__main__":
-    interface.launch()
+iface.launch()