AIの「根拠」を可視化せよ！ブラックボックスをこじ開けるCAMの正体

# 0. 画像をダウンロード
!wget -O cat.jpg https://saycon.co.jp/yuu/cat.jpg

import os
import numpy as np
import torch
from PIL import Image
import matplotlib.pyplot as plt
from torchvision import models
from torchvision.models import ResNet50_Weights
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
from pytorch_grad_cam.utils.image import show_cam_on_image

# 1. モデルの準備
weights = ResNet50_Weights.DEFAULT
model = models.resnet50(weights=weights)
model.eval()

target_layers = [model.layer4[-1]]

# 2. 画像の読み込みと前処理
input_path = "cat.jpg"
if not os.path.exists(input_path):
    raise FileNotFoundError(f"エラー: {input_path} が見つかりません。")

rgb_img = Image.open(input_path).convert("RGB")

# weights付属の標準前処理（Resize 256 -> CenterCrop 224 を含む）
preprocess = weights.transforms()

# 推論用テンソル作成（アスペクト比を維持したまま 224x224 に切り出される）
input_tensor = preprocess(rgb_img).unsqueeze(0)

# 【修正ポイント】可視化用に正規化を解除（Denormalize）
# weights.transforms 内にある mean と std を直接使用して計算します
mean = torch.tensor(preprocess.mean).view(3, 1, 1)
std = torch.tensor(preprocess.std).view(3, 1, 1)
img_for_vis = (input_tensor[0] * std + mean).clamp(0, 1).permute(1, 2, 0).numpy()

# 3. 予測結果からトップのクラスを取得
with torch.no_grad():
    outputs = model(input_tensor)

target_category = outputs.argmax(dim=1).item()
predicted_label = weights.meta["categories"][target_category]

print(f"AIの診断結果: {predicted_label} (Class ID: {target_category})")

# 4. Grad-CAMの実行
cam = GradCAM(model=model, target_layers=target_layers)
targets = [ClassifierOutputTarget(target_category)]

# cam() の戻り値は [batch, H, W] なので、最初の画像を取り出す
grayscale_cam = cam(input_tensor=input_tensor, targets=targets)[0, :]

# 5. 可視化
# img_for_vis (0-1 float) を背景にヒートマップを合成
visualization = show_cam_on_image(img_for_vis, grayscale_cam, use_rgb=True)

plt.figure(figsize=(12, 6))
plt.subplot(1, 2, 1)
plt.title(f"Original (Cropped): {predicted_label}")
plt.imshow(img_for_vis)
plt.axis("off")

plt.subplot(1, 2, 2)
plt.title(f"Grad-CAM Heatmap ({predicted_label})")
plt.imshow(visualization)
plt.axis("off")

plt.show()