Updated

- Blur the sensitive parts in the uploaded image
- Integrated Explainable AI

app.py

@@ -1,16 +1,20 @@
 import streamlit as st
 import torch
+import numpy as np
 from PIL import Image
 from torchvision import transforms
 import json
-from src.model import SwinTransformerMultiLabel  # Import from src folder
+import matplotlib.pyplot as plt
+import cv2
+from src.model import SwinTransformerMultiLabel
+from torchcam.methods import SmoothGradCAMpp  # Explainability Module
 
 # Title and description
-st.title("STAR Multi-Label Classifier")
-st.write("Upload an image to get multi-label predictions.")
+st.title("STAR Multi-Label Classifier with Sensitive Content Blurring")
+st.write("Upload an image to classify and see the blurred output.")
 
 # Load class labels from JSON
-label_file = "data/labels.json"  # Path to labels file
+label_file = "data/labels.json"
 with open(label_file, "r") as f:
     label_data = json.load(f)
 
@@ -24,6 +28,9 @@ model = SwinTransformerMultiLabel(num_classes=NUM_CLASSES)
 model.load_state_dict(torch.load(model_path, map_location="cpu"))
 model.eval()
 
+# Initialize CAM explainability
+cam_extractor = SmoothGradCAMpp(model, target_layer="model.features.7.3")  # Choose a valid layer
+
 # Define image preprocessing transformations
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
@@ -31,18 +38,41 @@ transform = transforms.Compose([
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
 ])
 
+# Function to blur sensitive areas
+def blur_sensitive_parts(image, cam_mask, blur_intensity=25):
+    image = np.array(image)  # Convert PIL to NumPy
+    heatmap = cv2.resize(cam_mask, (image.shape[1], image.shape[0]))  # Resize CAM to match image
+    heatmap = (heatmap > 0.6).astype(np.uint8)  # Threshold for blurring region
+
+    blurred = cv2.GaussianBlur(image, (51, 51), blur_intensity)  # Apply blur
+    mask = np.repeat(heatmap[:, :, np.newaxis], 3, axis=2)  # Create 3-channel mask
+    result = np.where(mask == 1, blurred, image)  # Blend blurred areas
+
+    return Image.fromarray(result)  # Convert back to PIL Image
+
 # Upload image
 uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
 if uploaded_file is not None:
     image = Image.open(uploaded_file).convert("RGB")
-    st.image(image, caption="Uploaded Image", use_column_width=True)
 
-    # Predict
+    # Preprocess image for model
     img_tensor = transform(image).unsqueeze(0)
+
+    # Perform inference
     with torch.no_grad():
         output = model(img_tensor)
-        predicted_indices = [i for i in range(NUM_CLASSES) if output[0][i] > 0.5]  # Threshold = 0.5
-        predicted_labels = [class_labels[i] for i in predicted_indices]  # Convert indices to labels
+        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
+    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
+    st.image(blurred_image, caption="Blurred Output Image", use_column_width=True)
 
-    # Display results
+    # Display predictions
     st.write("✅ **Predicted Labels:**", ", ".join(predicted_labels) if predicted_labels else "No labels detected")