Upload gradcam_xception.py

gradcam_xception.py

@@ -0,0 +1,460 @@
+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