Update app.py

app.py

@@ -1,148 +1,36 @@
-import torch
-import torch.nn as nn
-from torchvision import transforms as T
-import os
 import gradio as gr
-
-#################################
-# Define problem parameters
-#################################
-
-class config:
-
-    img_size = 224
-
-    pn_mean = [0.4752, 0.4752, 0.4752]  # Pneumonia dataset mean
-    pn_std = [0.2234, 0.2234, 0.2234]   # Pneumonia dataset std
-
-    class_names = ["Normal", "Pneumonia"]
-
-
-device = torch.device('cpu')
-
-print(f"device: {device}")
-
-#######################################
-# Define image transformation pipeline
-#######################################
-
-class Gray2RGB:
-    def __call__(self, image):
-        if image.shape[0] == 3:
-            return image
-        else:
-            return image.repeat(3, 1, 1) # Repeat the single channel across 3 channels to convert to RGB
-
-test_transform_custom = T.Compose([
-                            T.Resize(size=(config.img_size, config.img_size)),
-                            T.ToTensor(),
-                            Gray2RGB(),
-                            T.Normalize(config.pn_mean, config.pn_std),
-                        ])
-
-
-#################################
-# Define model architecture
-#################################
-
-class ConvolutionalNetwork(nn.Module):
-    def __init__(self):
-        super().__init__()
-
-        self.conv1 = nn.Sequential(
-                        nn.Conv2d(3, 8, 3, stride=1, padding=1),
-                        nn.ReLU(inplace=True),
-                        nn.BatchNorm2d(8),
-                        nn.MaxPool2d(2,2))
-
-        self.conv2 = nn.Sequential(
-                        nn.Conv2d(8, 16, 3, stride=1, padding=1),
-                        nn.ReLU(inplace=True),
-                        nn.BatchNorm2d(16),
-                        nn.MaxPool2d(2,2))
-
-        self.conv3 = nn.Sequential(
-                        nn.Conv2d(16, 32, 3, stride=1, padding=1),
-                        nn.ReLU(inplace=True),
-                        nn.BatchNorm2d(32),
-                        nn.MaxPool2d(2,2))
-
-        self.conv4 = nn.Sequential(
-                        nn.Conv2d(32, 64, 3, stride=1, padding=1),
-                        nn.ReLU(inplace=True),
-                        nn.BatchNorm2d(64),
-                        nn.MaxPool2d(2,2))
-
-        self.conv5 = nn.Sequential(
-                        nn.Conv2d(64, 128, 3, stride=1, padding=1),
-                        nn.ReLU(inplace=True),
-                        nn.BatchNorm2d(128),
-                        nn.MaxPool2d(2,2))
-
-        self.fc = nn.Sequential(
-                    nn.Linear(128*7*7, 512),
-                    nn.ReLU(inplace=True),
-                    nn.BatchNorm1d(512),
-                    nn.Dropout(0.5),
-                    nn.Linear(512, 2))
-
-    def forward(self, x):
-        x = self.conv1(x)
-        x = self.conv2(x)
-        x = self.conv3(x)
-        x = self.conv4(x)
-        x = self.conv5(x)
-        x = x.view(x.shape[0], -1)
-        x = self.fc(x)
-        return x
-
-
-cnn_model = ConvolutionalNetwork()
-
-cnn_model.to(device)
-
-status = cnn_model.load_state_dict(torch.load('pneumonia_cnn_model.pt', map_location=device, weights_only=True))
-print(f"Status: {status}")
-
-#################################
-# Define the prediction fucntion
-#################################
-
-def predict(image):
-    """Transforms and performs a prediction on an image and returns the prediction dictionary."""
-    
-    image = test_transform_custom(image).unsqueeze(0)
-    
-    cnn_model.eval()
-    with torch.no_grad():
-        pred_probs = torch.softmax(cnn_model(image), dim=1)
-    
-    # Create a prediction probability dictionary for each prediction class
-    pred_dict = {config.class_names[i]: float(pred_probs[0][i]) for i in range(len(config.class_names))}
-    
-    # Return the prediction dictionary
-    return pred_dict
-
-##########################
-# Create the Gradio demo
-##########################
-
-title = "Pneumonia Detection"
-
-description = """This is a pneumonia detection model that uses a custom convolutional neural network to predict whether an image contains pneumonia or not. \
-GitHub project can be accessed [here](https://github.com/mma666/Pneumonia-Detection-Computer-Vision).
-"""
-
-# Create examples list from "examples/" directory
-example_list = [["examples/" + example] for example in os.listdir("examples")]
-
-# Create the Gradio demo
-demo = gr.Interface(fn=predict,
-                    inputs=[gr.Image(label="Upload image", type="pil", height=320, width=320)],
-                    outputs=[gr.Label(num_top_classes=2, label="Predictions")],
-                    examples=example_list, 
-                    title=title,
-                    description=description,
-                    cache_examples=False)
-
-demo.launch()
+import numpy as np
+import tensorflow as tf
+from tensorflow.keras.models import load_model
+from tensorflow.keras.preprocessing import image
+
+# Load the trained model
+model = load_model("pneumonia_cnn_model.h5", custom_objects={
+    'f1': lambda y_true, y_pred: 2*((precision(y_true, y_pred)*recall(y_true, y_pred)) / (precision(y_true, y_pred)+recall(y_true, y_pred)+tf.keras.backend.epsilon())),
+    'precision': precision,
+    'recall': recall
+})
+
+# Preprocessing function
+def predict_pneumonia(img):
+    img = img.convert('L')  # convert to grayscale
+    img = img.resize((299, 299))
+    img_array = image.img_to_array(img)
+    img_array = img_array / 255.0
+    img_array = np.expand_dims(img_array, axis=0)  # batch dimension
+    prediction = model.predict(img_array)[0][0]
+
+    if prediction >= 0.5:
+        return "PNEUMONIA 🫁 (Confidence: {:.2f}%)".format(prediction * 100)
+    else:
+        return "NORMAL ✅ (Confidence: {:.2f}%)".format((1 - prediction) * 100)
+
+# Gradio interface
+interface = gr.Interface(fn=predict_pneumonia,
+                         inputs=gr.Image(type="pil"),
+                         outputs="text",
+                         title="Pneumonia Detection from Chest X-rays",
+                         description="Upload a chest X-ray image to detect Pneumonia using a CNN model.")
+
+if __name__ == "__main__":
+    interface.launch()