Update app.py

app.py

@@ -1,78 +1,50 @@
 import gradio as gr
-import torch
-import torch.nn as nn
-from torchvision import transforms
+import tensorflow as tf
 from PIL import Image
+import numpy as np
 import logging
 
 # Set up logging for debugging
 logging.basicConfig(level=logging.DEBUG)
 
-# Define the model architecture that matches the saved .pth file
-class BacterialMorphologyClassifier(nn.Module):
-    def __init__(self):
-        super(BacterialMorphologyClassifier, self).__init__()
-        self.feature_extractor = nn.Sequential(
-            nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1),
-            nn.ReLU(),
-            nn.MaxPool2d(kernel_size=2, stride=2),
-            nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),
-            nn.ReLU(),
-            nn.MaxPool2d(kernel_size=2, stride=2),
-        )
-        self.fc = nn.Sequential(
-            nn.Flatten(),
-            nn.Linear(64 * 56 * 56, 128),
-            nn.ReLU(),
-            nn.Dropout(0.5),
-            nn.Linear(128, 3),
-            nn.Softmax(dim=1),
-        )
-    
-    def forward(self, x):
-        x = self.feature_extractor(x)
-        x = self.fc(x)
-        return x
-
-# Load the model and weights
+# Load the .keras model
 MODEL_PATH = "https://huggingface.co/yolac/BacterialMorphologyClassification/resolve/main/model.keras"
 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()
+    model = tf.keras.models.load_model(MODEL_PATH)
     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([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-])
+def preprocess_image(image):
+    logging.debug("Preprocessing image...")
+    image = image.resize((224, 224))  # Resize to match model input size
+    image_array = np.array(image) / 255.0  # Normalize pixel values
+    if len(image_array.shape) == 2:  # If grayscale, convert to RGB
+        image_array = np.stack([image_array] * 3, axis=-1)
+    image_array = np.expand_dims(image_array, axis=0)  # Add batch dimension
+    return image_array
 
 # Define the prediction function
 def predict(image):
     try:
-        logging.debug("Starting prediction...")
         # Preprocess the image
-        image_tensor = transform(image).unsqueeze(0)
+        image_array = preprocess_image(image)
         logging.debug("Image preprocessing completed.")
         
         # Make prediction
-        output = model(image_tensor)
-        prediction = output.argmax().item()
+        predictions = model.predict(image_array)
+        prediction = np.argmax(predictions, axis=1)[0]
         
         # Class mapping
         class_labels = {0: 'cocci', 1: 'bacilli', 2: 'spirilla'}
         
         # Log prediction details
-        logging.debug(f"Predicted class: {class_labels[prediction]}, Confidence: {output.max().item()}")
+        logging.debug(f"Predicted class: {class_labels[prediction]}, Confidence: {predictions[0][prediction]}")
         
         # Return prediction result
-        return class_labels[prediction], float(output.max().item())
+        return class_labels[prediction], float(predictions[0][prediction])
     except Exception as e:
         logging.error(f"Error during prediction: {str(e)}")
         return "Error", 0.0