Update app.py

app.py

@@ -1,125 +1,57 @@
 import torch
-from transformers import ViTForImageClassification, ViTImageProcessor
+import torch.nn as nn
+from torchvision import transforms, models
 from PIL import Image
 import gradio as gr
-import numpy as np
-import matplotlib.pyplot as plt
 
-# Load the pretrained Vision Transformer model and image processor
-model_name = "google/vit-base-patch16-224"
-try:
-    model = ViTForImageClassification.from_pretrained(model_name)
-except Exception as e:
-    print(f"Error loading model: {e}")
-image_processor = ViTImageProcessor.from_pretrained(model_name)
+# Load the pre-trained DenseNet-121 model
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = models.densenet121(pretrained=True)
 
-# NIH Chest X-ray predefined conditions
-labels = [
-    "Atelectasis", "Cardiomegaly", "Effusion", "Infiltration", "Mass", "Nodule",
-    "Pneumonia", "Pneumothorax", "Consolidation", "Edema", "Emphysema", 
-    "Fibrosis", "Pleural Thickening", "Hernia"
-]
+# Modify the classifier layer to output probabilities for 14 classes (14 pathologies)
+num_classes = 14
+model.classifier = nn.Sequential(
+    nn.Linear(model.classifier.in_features, num_classes),
+    nn.Sigmoid()  # Use Sigmoid for multi-label classification
+)
+model.load_state_dict(torch.load('chexnet.pth', map_location=device))  # Load your pre-trained weights
+model = model.to(device)
+model.eval()
 
-# Function to apply Grad-CAM visualization
-def generate_grad_cam(image, target_layer):
-    try:
-        # Convert image to RGB if necessary
-        if image.mode != 'RGB':
-            image = image.convert('RGB')
-        
-        # Preprocess the image
-        inputs = image_processor(images=image, return_tensors="pt")
-        
-        # Forward pass to get logits
-        input_tensor = inputs["pixel_values"]
-        input_tensor.requires_grad = True  # Enable gradient tracking
-        outputs = model(input_tensor)
-        logits = outputs.logits
-        
-        # Get the predicted class and calculate gradients
-        predicted_class = logits.argmax(-1)
-        class_score = logits[0, predicted_class]
-        class_score.backward()
-        
-        # Get gradients and weights from the target layer
-        gradients = model.get_input_embeddings().weight.grad
-        pooled_gradients = torch.mean(gradients, dim=[0, 2, 3])
-        
-        # Apply Grad-CAM calculation (modify this part as per the model architecture)
-        cam = torch.mean(pooled_gradients * inputs["pixel_values"], dim=1).squeeze()
-        cam = torch.clamp(cam, min=0).numpy()  # Ensure non-negative values
-        
-        return cam, predicted_class.item(), None  # No error
-    except Exception as e:
-        error_message = f"Error generating Grad-CAM: {e}"
-        print(error_message)
-        return None, None, error_message
+# Define image transformations (resize, normalize)
+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])
+])
+
+# Class names for the 14 diseases (labels from ChestX-ray14 dataset)
+class_names = [
+    'Atelectasis', 'Cardiomegaly', 'Effusion', 'Infiltration', 'Mass', 
+    'Nodule', 'Pneumonia', 'Pneumothorax', 'Consolidation', 'Edema', 
+    'Emphysema', 'Fibrosis', 'Pleural Thickening', 'Hernia'
+]
 
-# Function to predict classes and visualize Grad-CAM
-def predict_and_explain(image):
-    try:
-        # Convert image to RGB if necessary
-        if image.mode != 'RGB':
-            image = image.convert('RGB')
-        
-        # Preprocess the image
-        inputs = image_processor(images=image, return_tensors="pt")
-        
-        # Forward pass to get logits
-        input_tensor = inputs["pixel_values"]
-        outputs = model(input_tensor)
-        logits = outputs.logits
-        
-        predicted_class = logits.argmax(-1).item()
-        cam_map, _, grad_cam_error = generate_grad_cam(image, "pooler_output")
-        
-        # Check for Grad-CAM errors
-        if grad_cam_error is not None:
-            return {
-                "predicted class": "Error during Grad-CAM generation",
-                "Grad-CAM map": None,
-                "error log": grad_cam_error
-            }
-        
-        # Convert cam_map to a visualizable format (heatmap)
-        if cam_map is not None:
-            plt.imshow(cam_map, cmap='jet', alpha=0.5)
-            plt.axis('off')
-            plt.title(f"Grad-CAM for {labels[predicted_class]}")
-            plt.colorbar()
-            plt.savefig("grad_cam_output.png")
-            plt.close()
-        
-            # Load the saved image to return it
-            grad_cam_image = Image.open("grad_cam_output.png")
-        else:
-            grad_cam_image = None
-        
-        return {
-            "predicted class": labels[predicted_class],
-            "Grad-CAM map": grad_cam_image,
-            "error log": "No errors"
-        }
-    except Exception as e:
-        error_message = f"Error predicting and explaining: {e}"
-        print(error_message)
-        return {
-            "predicted class": "Error during prediction",
-            "Grad-CAM map": None,
-            "error log": error_message
-        }
+# Prediction function
+def predict_disease(image):
+    image = transform(image).unsqueeze(0).to(device)  # Transform and add batch dimension
+    with torch.no_grad():
+        outputs = model(image)
+        outputs = outputs.cpu().numpy().flatten()
+    
+    # Create a dictionary of disease probabilities
+    result = {class_name: float(prob) for class_name, prob in zip(class_names, outputs)}
+    return result
 
-# Use the updated Gradio components syntax
-iface = gr.Interface(
-    fn=predict_and_explain,
-    inputs=gr.Image(type="pil"),  # Proper input type for images
-    outputs=[
-        gr.Textbox(label="Predicted Class"),
-        gr.Image(label="Grad-CAM Map"),
-        gr.Textbox(label="Error Log")
-    ],
-    title="Chest X-ray Classification with Debugging Logs"
+# Gradio Interface
+interface = gr.Interface(
+    fn=predict_disease,
+    inputs=gr.inputs.Image(type='pil'),  # Input is an image
+    outputs="label",  # Output is a dictionary of labels with probabilities
+    title="CheXNet Pneumonia Detection",
+    description="Upload a chest X-ray to detect the probability of 14 different diseases."
 )
 
+# Launch the Gradio app
 if __name__ == "__main__":
-    iface.launch()
+    interface.launch()