Update app.py

app.py

@@ -1,92 +1,92 @@
-import torch
-from transformers import ViTForImageClassification, ViTFeatureExtractor
-from PIL import Image
-import gradio as gr
-import numpy as np
-import matplotlib.pyplot as plt
-
-# Load the pretrained Vision Transformer model and feature extractor
-model_name = "google/vit-base-patch16-224"
-model = ViTForImageClassification.from_pretrained(model_name)
-feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)
-model.eval()
-
-# Function to apply Grad-CAM visualization
-def generate_grad_cam(image, target_layer):
-    # Preprocess the image
-    inputs = feature_extractor(images=image, return_tensors="pt")
-    input_tensor = inputs['pixel_values']
-    
-    # Forward pass to get logits
-    input_tensor.requires_grad = True
-    outputs = model(input_tensor)
-    
-    # Get the target score
-    score = outputs.logits[0].max()
-    
-    # Backpropagate to get gradients
-    model.zero_grad()
-    score.backward()
-    
-    # Get the gradients and activations from the target layer
-    gradients = model.get_input_embeddings().weight.grad
-    activations = model.get_input_embeddings().weight.data
-    
-    # Calculate Grad-CAM
-    pooled_gradients = torch.mean(gradients, dim=[0, 2, 3])
-    for i in range(activations.size(1)):
-        activations[:, i, :, :] *= pooled_gradients[i]
-    
-    heatmap = torch.mean(activations, dim=1).squeeze()
-    heatmap = np.maximum(heatmap.detach().numpy(), 0)
-    heatmap = heatmap / np.max(heatmap)
-    
-    return heatmap
-
-# Prediction and Grad-CAM function
-def predict_and_explain(image):
-    # Predict the class
-    inputs = feature_extractor(images=image, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model(**inputs)
-    
-    logits = outputs.logits
-    predicted_class_idx = logits.argmax(-1).item()
-    
-    # Predefined medical conditions (adjust based on your dataset)
-    labels = ["Class 1 - Normal", "Class 2 - Condition A", "Class 3 - Condition B"]
-    predicted_label = labels[predicted_class_idx]
-    
-    # Generate Grad-CAM heatmap
-    heatmap = generate_grad_cam(image, target_layer="vit.encoder.layer.11.output")
-    
-    # Visualize the heatmap on the original image
-    img = np.array(image)
-    heatmap_resized = np.array(Image.fromarray(heatmap).resize((img.shape[1], img.shape[0])))
-    
-    # Overlay heatmap on the original image
-    plt.imshow(img)
-    plt.imshow(heatmap_resized, cmap='jet', alpha=0.5)
-    plt.axis('off')
-    
-    # Save the overlayed image
-    plt.savefig("grad_cam_result.png")
-    
-    return predicted_label, "grad_cam_result.png"
-
-# Gradio interface
-interface = gr.Interface(
-    fn=predict_and_explain,
-    inputs=gr.inputs.Image(type="pil"),
-    outputs=[
-        "text",
-        gr.outputs.Image(type="file", label="Grad-CAM Visualization")
-    ],
-    title="Medical Image Analysis Tool with Explainability",
-    description="Upload an X-ray or MRI image to get a prediction for a medical condition with explainability through Grad-CAM.",
-    live=True
-)
-
-# Launch the app
-if __name__ == "__main__":
-    interface.launch()
+import torch
+from transformers import ViTForImageClassification, ViTImageProcessor
+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"
+model = ViTForImageClassification.from_pretrained(model_name)
+image_processor = ViTImageProcessor.from_pretrained(model_name)
+model.eval()
+
+# Function to apply Grad-CAM visualization
+def generate_grad_cam(image, target_layer):
+    # Preprocess the image
+    inputs = image_processor(images=image, return_tensors="pt")
+    input_tensor = inputs['pixel_values']
+    
+    # Forward pass to get logits
+    input_tensor.requires_grad = True
+    outputs = model(input_tensor)
+    
+    # Get the target score
+    score = outputs.logits[0].max()
+    
+    # Backpropagate to get gradients
+    model.zero_grad()
+    score.backward()
+    
+    # Get the gradients and activations from the target layer
+    gradients = model.get_input_embeddings().weight.grad
+    activations = model.get_input_embeddings().weight.data
+    
+    # Calculate Grad-CAM
+    pooled_gradients = torch.mean(gradients, dim=[0, 2, 3])
+    for i in range(activations.size(1)):
+        activations[:, i, :, :] *= pooled_gradients[i]
+    
+    heatmap = torch.mean(activations, dim=1).squeeze()
+    heatmap = np.maximum(heatmap.detach().numpy(), 0)
+    heatmap = heatmap / np.max(heatmap)
+    
+    return heatmap
+
+# Prediction and Grad-CAM function
+def predict_and_explain(image):
+    # Predict the class
+    inputs = image_processor(images=image, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+    
+    logits = outputs.logits
+    predicted_class_idx = logits.argmax(-1).item()
+    
+    # Predefined medical conditions (adjust based on your dataset)
+    labels = ["Class 1 - Normal", "Class 2 - Condition A", "Class 3 - Condition B"]
+    predicted_label = labels[predicted_class_idx]
+    
+    # Generate Grad-CAM heatmap
+    heatmap = generate_grad_cam(image, target_layer="vit.encoder.layer.11.output")
+    
+    # Visualize the heatmap on the original image
+    img = np.array(image)
+    heatmap_resized = np.array(Image.fromarray(heatmap).resize((img.shape[1], img.shape[0])))
+    
+    # Overlay heatmap on the original image
+    plt.imshow(img)
+    plt.imshow(heatmap_resized, cmap='jet', alpha=0.5)
+    plt.axis('off')
+    
+    # Save the overlayed image
+    plt.savefig("grad_cam_result.png")
+    
+    return predicted_label, "grad_cam_result.png"
+
+# Gradio interface
+interface = gr.Interface(
+    fn=predict_and_explain,
+    inputs=gr.Image(type="pil"),
+    outputs=[
+        "text",
+        gr.Image(type="file", label="Grad-CAM Visualization")
+    ],
+    title="Medical Image Analysis Tool with Explainability",
+    description="Upload an X-ray or MRI image to get a prediction for a medical condition with explainability through Grad-CAM.",
+    live=True
+)
+
+# Launch the app
+if __name__ == "__main__":
+    interface.launch()