gradcam_xception.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
from torchvision.transforms.functional import to_pil_image
import matplotlib.pyplot as plt
from PIL import Image
import numpy as np
import cv2
from timm import create_model
from huggingface_hub import hf_hub_download
import io
import warnings

warnings.filterwarnings("ignore", category=UserWarning)

# Advanced Grad-CAM Implementation
class AdvancedGradCAM:
    def __init__(self, model, target_layer, method="gradcam"):
        self.model = model
        self.target_layer = target_layer
        self.method = method
        self.gradients = None
        self.activations = None
        self.forward_hook_handle = None
        self.backward_hook_handle = None
        self._register_hooks()

    def _register_hooks(self):
        layer = dict([*self.model.named_modules()])[self.target_layer]

        def forward_hook(module, input, output):
            if isinstance(output, tuple):
                for item in output:
                    if isinstance(item, torch.Tensor):
                        self.activations = item.detach()
                        break
            else:
                self.activations = output.detach()

        def backward_hook(module, grad_in, grad_out):
            self.gradients = grad_out[0].detach()

        self.forward_hook_handle = layer.register_forward_hook(forward_hook)
        self.backward_hook_handle = layer.register_backward_hook(backward_hook)

    def remove_hooks(self):
        if self.forward_hook_handle:
            self.forward_hook_handle.remove()
        if self.backward_hook_handle:
            self.backward_hook_handle.remove()
        self.forward_hook_handle = None
        self.backward_hook_handle = None
        self.gradients = None
        self.activations = None

    def generate(self, input_tensor, class_idx, num_samples=5, stdev_spread=0.15):
        if self.forward_hook_handle is None or self.backward_hook_handle is None:
            self._register_hooks()

        self.model.zero_grad()

        try:
            input_tensor.requires_grad_(True)
            output = self.model(input_tensor)
            class_score = output[:, class_idx]
            class_score.backward()

            if self.gradients is None or self.activations is None:
                print(f"Warning: Gradients or activations are None for layer {self.target_layer}. Using fallback CAM.")
                h, w = input_tensor.shape[-2:]
                fallback_h, fallback_w = h // 16, w // 16
                return np.ones((fallback_h, fallback_w), dtype=np.float32) * 0.5

            if self.method == "gradcam":
                cam_result = self._standard_gradcam()
            else:
                raise ValueError(f"Unsupported CAM method: {self.method}")

            self.gradients = None
            self.activations = None
            input_tensor.requires_grad_(False)
            self.model.zero_grad()

            return cam_result

        except Exception as e:
            print(f"Error in AdvancedGradCAM.generate: {str(e)}")
            import traceback
            traceback.print_exc()
            h, w = input_tensor.shape[-2:]
            fallback_h, fallback_w = h // 16, w // 16
            return np.ones((fallback_h, fallback_w), dtype=np.float32) * 0.5

    def _standard_gradcam(self):
        gradients = self.gradients.cpu().numpy()
        activations = self.activations.cpu().numpy()

        if len(gradients.shape) != 4 or len(activations.shape) != 4:
             print(f"Warning: Unexpected shape in GradCAM++. Gradients: {gradients.shape}, Activations: {activations.shape}. Using fallback.")
             fallback_h, fallback_w = activations.shape[-2:] if len(activations.shape) >= 2 else (14, 14)
             return np.ones((fallback_h, fallback_w), dtype=np.float32) * 0.5

        grad_2 = gradients ** 2
        grad_3 = gradients ** 3
        epsilon = 1e-10
        alpha_denom = 2 * grad_2 + np.sum(activations * grad_3, axis=(2, 3), keepdims=True)
        alpha = grad_2 / (alpha_denom + epsilon)
        positive_activations_gradients = np.maximum(gradients, 0)
        weights = np.sum(alpha * positive_activations_gradients, axis=(2, 3))

        cam = np.zeros(activations.shape[2:], dtype=np.float32)
        for i, w in enumerate(weights[0]):
            cam += w * activations[0, i, :, :]

        cam = np.maximum(cam, 0)
        if np.max(cam) > 0:
            cam = cam / np.max(cam)
        return cam 

# Utility Functions
def overlay_cam_on_image(image, cam, face_box=None, alpha=0.5):
    img_np = np.array(image.convert('RGB'))
    h, w = img_np.shape[:2]

    if face_box is not None:
        x, y, fw, fh = map(int, face_box)
        if fw <= 0 or fh <= 0:
            print(f"Warning: Invalid face box dimensions {fw}x{fh}. Applying CAM to full image.")
            face_box = None
        else:
            try:
                face_cam_resized = cv2.resize(cam, (fw, fh))
            except cv2.error as e:
                print(f"Error resizing CAM to face box {fw}x{fh}: {e}. Applying CAM to full image.")
                face_box = None

    if face_box is not None:
        x, y, fw, fh = map(int, face_box)
        full_cam_heatmap = np.zeros((h, w), dtype=np.float32)
        y_end = min(y + fh, h)
        x_end = min(x + fw, w)
        fh_clipped = y_end - y
        fw_clipped = x_end - x
        if fh_clipped > 0 and fw_clipped > 0:
            full_cam_heatmap[y:y_end, x:x_end] = face_cam_resized[:fh_clipped, :fw_clipped]
        else:
            print("Warning: Face box calculation resulted in zero area for heatmap placement.")
        heatmap_colored = cv2.applyColorMap(np.uint8(255 * full_cam_heatmap), cv2.COLORMAP_JET)
        heatmap_colored = cv2.cvtColor(heatmap_colored, cv2.COLOR_BGR2RGB)
    else:
        try:
            cam_resized = cv2.resize(cam, (w, h))
        except cv2.error as e:
            print(f"Error resizing CAM to full image size {w}x{h}: {e}. Skipping overlay.")
            return image
        heatmap_colored = cv2.applyColorMap(np.uint8(255 * cam_resized), cv2.COLORMAP_JET)
        heatmap_colored = cv2.cvtColor(heatmap_colored, cv2.COLOR_BGR2RGB)

    overlayed_img = cv2.addWeighted(img_np, 1 - alpha, heatmap_colored, alpha, 0)
    return Image.fromarray(overlayed_img)

def save_comparison(image, cam, overlay, face_box=None):
    fig, axes = plt.subplots(1, 3, figsize=(15, 5))

    axes[0].imshow(image)
    axes[0].set_title("Original")
    if face_box is not None:
        x, y, w, h = map(int, face_box)
        if w > 0 and h > 0:
            rect = plt.Rectangle((x, y), w, h, edgecolor='lime', linewidth=2, fill=False)
            axes[0].add_patch(rect)
    axes[0].axis("off")

    if face_box is not None:
        x, y, w, h = map(int, face_box)
        if w > 0 and h > 0:
            try:
                cam_display = cv2.resize(cam, (w, h))
                img_h, img_w = np.array(image).shape[:2]
                full_cam_display = np.zeros((img_h, img_w))
                y_end = min(y + h, img_h)
                x_end = min(x + w, img_w)
                h_clipped = y_end - y
                w_clipped = x_end - x
                if h_clipped > 0 and w_clipped > 0:
                    full_cam_display[y:y_end, x:x_end] = cam_display[:h_clipped, :w_clipped]
                axes[1].imshow(full_cam_display, cmap="jet")
            except cv2.error:
                axes[1].imshow(cam, cmap="jet")
        else:
            axes[1].imshow(cam, cmap="jet")
    else:
        axes[1].imshow(cam, cmap="jet")

    axes[1].set_title("CAM")
    axes[1].axis("off")

    axes[2].imshow(overlay)
    axes[2].set_title("Overlay")
    axes[2].axis("off")

    plt.tight_layout()
    
    buf = io.BytesIO()
    plt.savefig(buf, format="png", bbox_inches="tight")
    plt.close()
    buf.seek(0)
    return Image.open(buf)

#load xception model
def load_xception_model(model_repo="drg31/xception", model_filename="final_xception_model.pth", num_classes=2):
    try:
        model_path = hf_hub_download(repo_id=model_repo, filename=model_filename)
        print(f"Model downloaded to: {model_path}")
    except Exception as e:
        print(f"Error downloading model from Hugging Face Hub ({model_repo}/{model_filename}): {e}")
        raise

    model = create_model("xception", pretrained=False, num_classes=num_classes)
    print(f"Created Xception model with {num_classes} output classes.")

    try:
        checkpoint = torch.load(model_path, map_location=torch.device('cpu'), weights_only=False)
        print(f"Checkpoint loaded successfully from {model_path} (with weights_only=False).")
    except Exception as e:
        print(f"Error loading checkpoint from {model_path}: {e}")
        raise

    if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
        checkpoint_state_dict = checkpoint['state_dict']
        print("Extracted state_dict from checkpoint dictionary.")
    else:
        checkpoint_state_dict = checkpoint
        print("Using checkpoint directly as state_dict.")

    cleaned_state_dict = {}
    for k, v in checkpoint_state_dict.items():
        name = k.replace('module.', '')
        cleaned_state_dict[name] = v
    print(f"Cleaned state_dict contains {len(cleaned_state_dict)} keys (after removing 'module.' prefix).")

    print("Loading state_dict with strict=False...")
    report = model.load_state_dict(cleaned_state_dict, strict=False)
    print(f"Load report - Missing keys: {report.missing_keys}")
    print(f"Load report - Unexpected keys: {report.unexpected_keys}")

    print("Model state loaded.")
    model.eval()
    return model

def get_target_layer_xception(model):
    target_layer_name = "block12.rep.6"  # Deeper layer for semantic features
    if target_layer_name not in dict([*model.named_modules()]):
        print(f"Warning: Target layer '{target_layer_name}' not found. Trying 'block11.rep.2'.")
        target_layer_name = "block11.rep.2"
        if target_layer_name not in dict([*model.named_modules()]):
            print(f"Warning: Fallback layer '{target_layer_name}' not found. Trying 'act4'.")
            target_layer_name = "act4"
            if target_layer_name not in dict([*model.named_modules()]):
                raise ValueError("Could not find suitable target layer for GradCAM in Xception model.")
    print(f"Using target layer: {target_layer_name}")
    return target_layer_name

# Main Visualization Function
def generate_smoothgrad_visualizations_xception(model, image, target_class=None, face_only=True, num_samples=5, stdev_spread=0.15):
    print("\n--- Starting Prediction and Grad-CAM ---")
    try:
        predicted_class_idx, confidence = predict_image(model, image, face_only)
    except Exception as pred_e:
        print(f"Error during prediction: {pred_e}")
        import traceback
        traceback.print_exc()
        return None, None, None, None

    if target_class is None:
        cam_target_class = predicted_class_idx
        print(f"CAM Target Class: Using predicted class index {cam_target_class} ({'real' if cam_target_class == 0 else 'fake'})")
    elif target_class in [0, 1]:
        cam_target_class = target_class
        print(f"CAM Target Class: Using specified class index {cam_target_class} ({'real' if cam_target_class == 0 else 'fake'})")
    else:
        print(f"Warning: Invalid target_class specified ({target_class}). Defaulting to predicted class index {predicted_class_idx}.")
        cam_target_class = predicted_class_idx

    device = next(model.parameters()).device
    model.eval()
    
    IMAGE_SIZE = 299
    transform = transforms.Compose([
        transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
    ])
    
    dataset = ImageDataset(image, transform=transform, face_only=face_only)
    input_tensor, original_image, face_box = dataset[0]
    input_tensor = input_tensor.unsqueeze(0).to(device)
    print(f"Input tensor for CAM shape: {input_tensor.shape}, Face box: {face_box}")

    raw_cam = None
    try:
        target_layer = get_target_layer_xception(model)
        print(f"Using target layer for CAM: {target_layer}")
        cam_extractor = AdvancedGradCAM(model, target_layer, method="gradcam")
        raw_cam = cam_extractor.generate(input_tensor, cam_target_class, num_samples=num_samples, stdev_spread=stdev_spread)
    except Exception as cam_e:
        print(f"Error during CAM generation: {cam_e}")
        import traceback
        traceback.print_exc()

    cam_heatmap_img, overlay_img, comparison_img = None, None, None
    if raw_cam is None or not isinstance(raw_cam, np.ndarray) or raw_cam.size == 0:
        print("CAM generation failed or produced invalid result. Skipping visualization.")
    else:
        try:
            print("Generating visualizations...")
            img_h, img_w = np.array(original_image).shape[:2]
            heatmap_display_np = np.zeros((img_h, img_w), dtype=np.float32)
            if face_box:
                x, y, w_fb, h_fb = map(int, face_box)
                if w_fb > 0 and h_fb > 0:
                    cam_resized_face = cv2.resize(raw_cam, (w_fb, h_fb), interpolation=cv2.INTER_LINEAR)
                    y_end, x_end = min(y + h_fb, img_h), min(x + w_fb, img_w)
                    h_clip, w_clip = y_end - y, x_end - x
                    if h_clip > 0 and w_clip > 0:
                        heatmap_display_np[y:y_end, x:x_end] = cam_resized_face[:h_clip, :w_clip]
                else:
                    heatmap_display_np = cv2.resize(raw_cam, (img_w, img_h), interpolation=cv2.INTER_LINEAR)
            else:
                heatmap_display_np = cv2.resize(raw_cam, (img_w, img_h), interpolation=cv2.INTER_LINEAR)

            min_h, max_h = np.min(heatmap_display_np), np.max(heatmap_display_np)
            if max_h > min_h:
                heatmap_norm = (heatmap_display_np - min_h) / (max_h - min_h)
            else:
                heatmap_norm = np.zeros_like(heatmap_display_np)
            heatmap_rgb = (plt.cm.jet(heatmap_norm)[:, :, :3] * 255).astype(np.uint8)
            cam_heatmap_img = Image.fromarray(heatmap_rgb)
            print("  Heatmap generated.")

            overlay_img = overlay_cam_on_image(original_image, raw_cam, face_box)
            print("  Overlay generated.")

            if overlay_img:
                comparison_img = save_comparison(original_image, raw_cam, overlay_img, face_box)
                print("  Comparison generated.")
            else:
                print("  Skipping comparison image because overlay failed.")

        except Exception as vis_e:
            print(f"Error during visualization generation: {vis_e}")
            import traceback
            traceback.print_exc()

    print("--- Prediction and Grad-CAM Finished ---")
    return raw_cam, cam_heatmap_img, overlay_img, comparison_img

# Face Detection Dataset
class ImageDataset(torch.utils.data.Dataset):
    def __init__(self, image, transform=None, face_only=True):
        self.image = image
        self.transform = transform
        self.face_only = face_only
        try:
            self.face_detector = cv2.CascadeClassifier(
                cv2.data.haarcascades + 'haarcascade_frontalface_default.xml'
            )
            if self.face_detector.empty():
                raise IOError("Failed to load cascade file")
        except Exception as e:
            print(f"Error loading Haar Cascade: {e}. Face detection might fail.")
            class DummyDetector:
                def detectMultiScale(self, *args, **kwargs): return []
            self.face_detector = DummyDetector()

    def __len__(self):
        return 1

    def detect_face(self, image_np):
        gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
        faces = self.face_detector.detectMultiScale(gray, 1.1, 5)
        if len(faces) == 0:
            print("No face detected, using full image as fallback.")
            h, w = image_np.shape[:2]
            return (0, 0, w, h), image_np
        areas = [w * h for (x, y, w, h) in faces]
        idx = np.argmax(areas)  # Select largest face
        x, y, w, h = faces[idx]
        pad_x, pad_y = int(w * 0.05), int(h * 0.05)
        x1, y1 = max(0, x - pad_x), max(0, y - pad_y)
        x2, y2 = min(image_np.shape[1], x + w + pad_x), min(image_np.shape[0], y + h + pad_y)
        face_img = image_np[y1:y2, x1:x2]
        return (x1, y1, x2 - x1, y2 - y1), face_img

    def __getitem__(self, idx):
        image_np = np.array(self.image)
        original_image = self.image.copy()
        face_box_final = None
        processed_image = original_image

        if self.face_only:
            try:
                face_box, face_img_np = self.detect_face(image_np)
                if face_img_np.size == 0 or face_box[2] <= 0 or face_box[3] <= 0:
                    print("Warning: Face detection returned empty or invalid region. Using full image.")
                    face_box_final = None
                    processed_image = original_image
                else:
                    processed_image = Image.fromarray(face_img_np)
                    face_box_final = face_box
            except Exception as e:
                print(f"Error during face detection: {e}. Using full image.")
                face_box_final = None
                processed_image = original_image
        else:
            face_box_final = None
            processed_image = original_image

        if self.transform:
            tensor = self.transform(processed_image)
        else:
            tensor = transforms.ToTensor()(processed_image)

        return tensor, original_image, face_box_final

def predict_image(model, image, face_only=True):
    device = next(model.parameters()).device
    model.eval()

    IMAGE_SIZE = 299
    transform = transforms.Compose([
        transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
    ])

    dataset = ImageDataset(image, transform=transform, face_only=face_only)
    input_tensor, _, _ = dataset[0]
    input_tensor = input_tensor.unsqueeze(0).to(device)

    with torch.no_grad():
        logits = model(input_tensor)
        probabilities = F.softmax(logits, dim=1)

    pred_prob = probabilities[0].max().item()
    pred_class_idx = probabilities[0].argmax().item()
    pred_label = "real" if pred_class_idx == 0 else "fake"

    if pred_prob < 0.7:  # Example threshold
        print(f"Warning: Low confidence ({pred_prob:.4f}) detected. Model may need fine-tuning.")

    print(f"--- Prediction ---")
    print(f"Input Tensor Shape: {input_tensor.shape}")
    print(f"Logits: {logits.cpu().numpy()}")
    print(f"Probabilities: {probabilities.cpu().numpy()}")
    print(f"Predicted Class: {pred_label} (Index: {pred_class_idx})")
    print(f"Confidence: {pred_prob:.4f}")
    print(f"--------------------")

    return pred_class_idx, pred_prob