Update app.py

app.py

@@ -7,10 +7,11 @@ 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)
+try:
+    model = ViTForImageClassification.from_pretrained(model_name)
+except Exception as e:
+    print(f"Error loading model: {e}")
 image_processor = ViTImageProcessor.from_pretrained(model_name)
-model.eval()
-
 # NIH Chest X-ray predefined conditions
 labels = [
     "Atelectasis",
@@ -28,82 +29,54 @@ labels = [
     "Pleural Thickening",
     "Hernia"
 ]
-
 # 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
+    try:
+        # 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
+        
+        # Calculate Grad-CAM
+        cam_weights = torch.mean(torch.relu(logits), dim=(2, 3))
+        cam_map = (torch.unsqueeze(cam_weights, 1) * 
+                   torch.sigmoid(outputs.pooler_output)).sum(dim=1).squeeze()
+        
+        return cam_map.numpy(), logits.argmax(-1)
+    except Exception as e:
+        print(f"Error generating Grad-CAM: {e}")
+        return None
+# Function to predict classes and visualize Grad-CAM
 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()
-    
-    # Get the predicted label based on NIH Chest X-ray conditions
-    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 NIH Chest X-ray",
-    description="Upload a Chest X-ray image to get a prediction for common thoracic conditions based on the NIH dataset, with explainability through Grad-CAM.",
-    live=True
+    try:
+        # 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, _ = generate_grad_cam(image, "pooler_output")
+        
+        return {
+            "predicted class": labels[predicted_class],
+            "Grad-CAM map": cam_map,
+        }
+    except Exception as e:
+        print(f"Error predicting and explaining: {e}")
+        return None
+# Create a Gradio interface
+iface = gr.Interface(
+    fn=predict_and_explain, 
+    inputs="image", 
+    outputs=["text", "image"], 
+    title="Chest X-ray Classification"
 )
 
-# Launch the app
 if __name__ == "__main__":
-    interface.launch()
+    iface.launch()