import torch
import torch.nn as nn
from torchvision import models, transforms
import cv2
import numpy as np
from PIL import Image 

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

def get_all_conv_layers(model):
    return [module for module in model.modules() if isinstance(module, nn.Conv2d)]

def grad_cam_all_layers(image_path):
    # 图像预处理
    image = Image.open(image_path).convert('RGB')
    original_image = np.array(image)
    
    preprocess = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    
    input_tensor = preprocess(image).unsqueeze(0).to(device)
    original_image = cv2.resize(original_image, (224, 224))
    
    # 加载模型
    model = models.resnet18(pretrained=True).to(device)
    model.eval()
    
    # 获取所有卷积层
    conv_layers = get_all_conv_layers(model)
    print(f"Found {len(conv_layers)} convolutional layers")
    
    # 创建存储结构
    activations = {}
    gradients = {}
    
    # 注册钩子的函数
    def save_activation(layer_name):
        def hook(module, input, output):
            activations[layer_name] = output.detach().cpu()
        return hook
    
    def save_gradient(layer_name):
        def hook(module, grad_input, grad_output):
            gradients[layer_name] = grad_output[0].detach().cpu()
        return hook
    
    # 为所有卷积层注册钩子
    hooks = []
    for idx, layer in enumerate(conv_layers):
        layer_name = f"conv{idx+1}"
        hooks.append(layer.register_forward_hook(save_activation(layer_name)))
        hooks.append(layer.register_backward_hook(save_gradient(layer_name)))
    
    # 前向传播
    output = model(input_tensor)
    pred_class = output.argmax(dim=1).item()
    
    # 反向传播
    model.zero_grad()
    one_hot = torch.zeros_like(output)
    print(output.shape)
    one_hot[0][pred_class] = 1
    output.backward(gradient=one_hot)
    

    # 为每个卷积层生成CAM
    for i, (layer_name, layer) in enumerate(zip(activations.keys(), conv_layers)):
        # 获取对应层的激活和梯度
        activation = activations[layer_name].squeeze()
        gradient = gradients[layer_name].squeeze()
        
        # 计算权重
        weights = gradient.mean(dim=(1, 2), keepdim=True)
        
        # 生成CAM
        cam = (weights * activation).sum(dim=0)
        cam = torch.relu(cam)
        cam = (cam - cam.min()) / (cam.max() - cam.min())
        cam = cam.numpy()
        
        # 调整尺寸
        cam = cv2.resize(cam, (224, 224))
        cam = np.uint8(255 * cam)
        heatmap = cv2.applyColorMap(cam, cv2.COLORMAP_JET)
        
        # 叠加图像
        superimposed_img = cv2.addWeighted(original_image, 0.6, heatmap, 0.4, 0)
        
        #use cv2 save grad cam
        cv2.imwrite(f'./cam/grad_cam_{layer_name}.jpg', superimposed_img)
           
    # 移除所有钩子
    for hook in hooks:
        hook.remove()

# 使用示例
grad_cam_all_layers('1.jpg')