Updated

- Error in XAI

app.py

@@ -7,7 +7,7 @@ import json
 import matplotlib.pyplot as plt
 import cv2
 from src.model import SwinTransformerMultiLabel
-from torchcam.methods import SmoothGradCAMpp  # Explainability Module
+from torchcam.methods import SmoothGradCAMpp  
 
 # Title and description
 st.title("STAR Multi-Label Classifier with Sensitive Content Blurring")
@@ -28,8 +28,21 @@ model = SwinTransformerMultiLabel(num_classes=NUM_CLASSES)
 model.load_state_dict(torch.load(model_path, map_location="cpu"))
 model.eval()
 
+# Print model architecture to find a valid layer for CAM
+st.write("🔍 Finding a valid layer for CAM...")
+valid_layer = None
+for name, module in model.named_modules():
+    if "conv" in name.lower() or "features" in name:  # Look for convolutional layers
+        valid_layer = name
+
+# Ensure a valid layer was found
+if not valid_layer:
+    raise ValueError("❌ No valid convolutional layer found for CAM!")
+
+st.write(f"✅ Using layer '{valid_layer}' for explainability.")
+
 # Initialize CAM explainability
-cam_extractor = SmoothGradCAMpp(model, target_layer="model.features.7.3")  # Choose a valid layer
+cam_extractor = SmoothGradCAMpp(model, target_layer=valid_layer)  # Use detected layer
 
 # Define image preprocessing transformations
 transform = transforms.Compose([
@@ -64,14 +77,14 @@ if uploaded_file is not None:
         predicted_indices = [i for i in range(NUM_CLASSES) if output[0][i] > 0.5]
         predicted_labels = [class_labels[i] for i in predicted_indices]
 
-    # Generate CAM heatmap
+    # Generate CAM heatmap and blur the image
     blurred_image = image  # Default to original if no prediction
     if predicted_indices:
         cam = cam_extractor(predicted_indices[0], output).squeeze().cpu().numpy()
         cam = (cam - np.min(cam)) / (np.max(cam))  # Normalize
         blurred_image = blur_sensitive_parts(image, cam)  # Apply blur
 
-    # Display the blurred image instead of the original
+    # Display only the blurred image
     st.image(blurred_image, caption="Blurred Output Image", use_column_width=True)
 
     # Display predictions