final_project_Gradio_App.py

# Let us first import all the necessary libraries required for this project

import tensorflow as tf
import torch
import cv2
import sklearn
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import os
from torchvision import datasets, transforms, models
from torch.utils.data import DataLoader, Dataset
import torch.nn as nn
import torch.optim as optim
import torchvision
import torch.nn.functional as F
from PIL import Image

from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score, confusion_matrix, roc_curve, auc


# Later on, as per requirement, more libraries wil be imported

import gradio as gr

def adjust_brightness_contrast(image, alpha=1.2, beta=50):
    """
    Adjusting brightness and contrast of the image.
    Parameters:
        - image: Input image (numpy array).
        - alpha: Contrast control [1.0-3.0].
        - beta: Brightness control [0-100].
    Returns:
        - Adjusted image.
    """
    return cv2.convertScaleAbs(image, alpha=alpha, beta=beta)

def apply_histogram_equalization(image):
    """Applying histogram equalization to enhance contrast."""
    channels = cv2.split(image)
    eq_channels = [cv2.equalizeHist(ch) for ch in channels]
    return cv2.merge(eq_channels)

def apply_clahe(image, clip_limit=2.0, tile_grid_size=(8, 8)):
    """Applying CLAHE for local contrast enhancement."""
    clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=tile_grid_size)
    channels = cv2.split(image)
    clahe_channels = [clahe.apply(ch) for ch in channels]
    return cv2.merge(clahe_channels)

def apply_gaussian_blur(image, kernel_size=(3, 3)):
    """Applying Gaussian blur for denoising."""
    return cv2.GaussianBlur(image, kernel_size, 0)

def apply_sharpening(image):
    """Applying edge enhancement using a sharpening filter."""
    kernel = np.array([[0, -1, 0],
                       [-1, 5, -1],
                       [0, -1, 0]])
    return cv2.filter2D(image, -1, kernel)

def normalize_image(image):
    """Normalizing the image to zero mean and unit variance."""
    image = (image - np.mean(image)) / np.std(image)
    return image

def resize_image(image, width, height):
    """Resizing the image to the desired dimensions with anti-aliasing."""
    return cv2.resize(image, (width, height), interpolation=cv2.INTER_CUBIC)


def preprocess_single_image(pil_image, img_height=224, img_width=224):
    """
    Preprocessing a single image as per the training pipeline.
    Parameters:
        - pil_image: Input PIL image.
        - img_height, img_width: Dimensions to resize the image.
    Returns:
        - Preprocessed image tensor.
    """
    # Converting PIL image to numpy array
    image = np.array(pil_image)

    # Ensuring the image is in RGB format
    if len(image.shape) == 2 or image.shape[2] == 1:  # Grayscale image
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)

    # Applying preprocessing steps
    image = apply_histogram_equalization(image)
    image = apply_clahe(image)
    image = apply_gaussian_blur(image)
    image = apply_sharpening(image)
    image = adjust_brightness_contrast(image, alpha=1.2, beta=50)

    # Resizing and normalization
    image = resize_image(image, img_width, img_height)
    image = normalize_image(image)

    # Converting to PIL image and applying transformations
    image = Image.fromarray(image.astype(np.uint8))
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    image_tensor = transform(image).unsqueeze(0)  # Adding batch dimension
    return image_tensor


# Detecting GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

from torchvision.models import densenet121
from torchvision.transforms import Resize

# Simplified ViT-like transformer module
class SimpleViT(nn.Module):
    def __init__(self, input_dim, num_heads, mlp_dim, num_layers):
        super(SimpleViT, self).__init__()
        # Reduced TransformerEncoder layer complexity
        self.transformer_blocks = nn.ModuleList([
            nn.TransformerEncoderLayer(
                d_model=input_dim,
                nhead=num_heads,
                dim_feedforward=mlp_dim,
                dropout=0.1
            ) for _ in range(num_layers)
        ])

    def forward(self, x):
        # Flattening the spatial dimensions
        B, C, H, W = x.shape
        x = x.flatten(2).permute(2, 0, 1)  # Reshaping for transformer
        for block in self.transformer_blocks:
            x = block(x)
        x = x.permute(1, 2, 0).reshape(B, C, H, W)  # Restoring the original shape
        return x


# Adjusted Hybrid DenseNet + Simplified ViT Architecture
class LightweightHybridDenseNetViT(nn.Module):
    def __init__(self):
        super(LightweightHybridDenseNetViT, self).__init__()

        # Loading a lighter DenseNet backbone
        self.densenet = densenet121(pretrained=False)  # Base DenseNet backbone

        # Reducing the output channels from DenseNet to smaller dimensions
        self.conv_reduce = nn.Conv2d(1024, 64, kernel_size=1)  # Fewer channels

        # ViT processing module with reduced complexity
        self.vit = SimpleViT(input_dim=64, num_heads=2, mlp_dim=128, num_layers=1)

        # Task-specific classification heads
        self.fc_pneumonia = nn.Linear(64, 1)  # Binary classification (Pneumonia)
        self.fc_tuberculosis = nn.Linear(64, 1)  # Binary classification (Tuberculosis)
        self.fc_lung_cancer = nn.Linear(64, 4)  # Multi-class output (Lung Cancer)

    def forward(self, x):
        # Extracting DenseNet features
        x = self.densenet.features(x)  # Extracting DenseNet feature maps
        x = self.conv_reduce(x)  # Reducing the number of feature channels

        # Passing through simplified ViT module
        x = self.vit(x)

        # Applying Global Average Pooling (GAP)
        x = x.mean(dim=[2, 3])  # Pooling across spatial dimensions

        # Task-specific classification
        pneumonia_output = torch.sigmoid(self.fc_pneumonia(x))  # Binary sigmoid output
        tuberculosis_output = torch.sigmoid(self.fc_tuberculosis(x))  # Binary sigmoid output
        lung_cancer_output = self.fc_lung_cancer(x)  # Multi-class logits

        return pneumonia_output, tuberculosis_output, lung_cancer_output
    

# Loading the saved model
img_size = 224  # Matching the dimensions used during training
patch_size = 8


model = LightweightHybridDenseNetViT().to(device)

model.load_state_dict(torch.load("model_FINAL.pth", map_location=device))  # Mapping to the correct device
model.to(device)  # Moving the model to GPU/CPU
model.eval()  # Setting to evaluation mode


# Function to pre-process the image and perform inference
def predict_image(image):
    """
    Predicts the probabilities of Pneumonia, TB, and Lung Cancer from the input image.
    """
    # Preprocessing the image
    image_tensor = preprocess_single_image(image, img_height=224, img_width=224)
    image_tensor = image_tensor.to(device)
    
    # Performing inference
    with torch.no_grad():
        pneumonia_output, tb_output, lung_cancer_output = model(image_tensor)
    
    # Getting the probabilities
    pneumonia_prob = pneumonia_output.item()
    tb_prob = tb_output.item()
    lung_cancer_probs = F.softmax(lung_cancer_output, dim=1).squeeze().tolist()

    # Class names for lung cancer
    lung_cancer_classes = [
        "adenocarcinoma_left.lower.lobe",
        "large.cell.carcinoma_left.hilum",
        "NORMAL",
        "squamous.cell.carcinoma_left.hilum"
    ]
    
    # Preparing the result as a dictionary
    result = {
        "Pneumonia Probability": f"{pneumonia_prob:.4f}",
        "TB Probability": f"{tb_prob:.4f}",
        "Lung Cancer Probabilities": {class_name: f"{prob:.4f}" for class_name, prob in zip(lung_cancer_classes, lung_cancer_probs)}
    }
    
    return result

# Gradio Interface
iface = gr.Interface(fn=predict_image, 
                     inputs=gr.Image(type="pil"), 
                     outputs=gr.JSON(),
                     title="Probabilistic Lung Disease Detection",
                     description="An AI-powered tool that analyzes and predicts probabilities for lung diseases, including Pneumonia, Tuberculosis, and Lung Cancer.")

iface.launch()