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