generate_gradcam_locally.py

#!/usr/bin/env python3
"""
Generate Grad-CAM visualizations for the thyroid model.
Run this locally with a GPU for best performance, or on CPU (slower).

Usage:
    python generate_gradcam_locally.py

Requirements:
    pip install torch torchvision transformers datasets matplotlib Pillow huggingface_hub
"""
import os, math, json, random, warnings, traceback
warnings.filterwarnings("ignore")

import numpy as np
from PIL import Image
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

import torch
import torch.nn.functional as F
from datasets import load_dataset
from transformers import AutoImageProcessor, AutoModelForImageClassification
from huggingface_hub import HfApi

# ============== CONFIG ==============
HF_USERNAME = "Johnyquest7"
DATASET_NAME = "BTX24/thyroid-cancer-classification-ultrasound-dataset"
MODEL_NAME = f"{HF_USERNAME}/ML-Inter_thyroid"
OUTPUT_DIR = "./gradcam_outputs"
REPO_ID = f"{HF_USERNAME}/thyroid-training-scripts"
SEED = 42
BATCH_SIZE = 16
NUM_CORRECT = 5   # Number of correct predictions to visualize
NUM_WRONG = 5     # Number of incorrect predictions to visualize
# =====================================

random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)

def main():
    print("=" * 60)
    print("Thyroid Grad-CAM Visualization Generator")
    print("=" * 60)
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"\nDevice: {device}")
    if device.type == "cpu":
        print("WARNING: Running on CPU. This will be slow for SwinV2 backward passes.")
        print("Consider running on Google Colab or a machine with GPU.")

    print(f"Loading model: {MODEL_NAME}")
    processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
    model = AutoModelForImageClassification.from_pretrained(MODEL_NAME).to(device).eval()
    id2label = model.config.id2label
    print(f"Model loaded: {sum(p.numel() for p in model.parameters())/1e6:.1f}M params")

    print(f"\nLoading dataset: {DATASET_NAME}")
    ds = load_dataset(DATASET_NAME, split="train")
    ds = ds.shuffle(seed=SEED)
    train_test = ds.train_test_split(test_size=0.2, stratify_by_column="label", seed=SEED)
    test_ds = train_test["test"]
    print(f"Test samples: {len(test_ds)}")

    # Get predictions
    all_logits, all_labels = [], []
    print("\nRunning inference...")
    for i in range(0, len(test_ds), BATCH_SIZE):
        batch_items = [test_ds[j] for j in range(i, min(i+BATCH_SIZE, len(test_ds)))]
        images = [item["image"].convert("RGB") for item in batch_items]
        inputs = processor(images, return_tensors="pt")
        with torch.no_grad():
            outputs = model(pixel_values=inputs["pixel_values"].to(device))
        all_logits.extend(outputs.logits.cpu().numpy())
        all_labels.extend([item["label"] for item in batch_items])
        if (i // BATCH_SIZE) % 5 == 0:
            print(f"  Batch {i//BATCH_SIZE + 1}/{(len(test_ds)+BATCH_SIZE-1)//BATCH_SIZE}")

    y_true = np.array(all_labels)
    y_pred = np.argmax(np.array(all_logits), axis=1)

    correct_idx = [i for i in range(len(y_true)) if y_true[i] == y_pred[i]]
    incorrect_idx = [i for i in range(len(y_true)) if y_true[i] != y_pred[i]]
    random.shuffle(correct_idx)
    random.shuffle(incorrect_idx)
    selected = correct_idx[:NUM_CORRECT] + incorrect_idx[:NUM_WRONG]
    print(f"\nSelected {len(selected)} samples: {len(correct_idx[:NUM_CORRECT])} correct, {len(incorrect_idx[:NUM_WRONG])} incorrect")

    # Register hooks on last SwinV2 stage
    gradcam_data = {}
    def fwd_hook(module, input, output):
        gradcam_data["feat"] = output.detach()
    def bwd_hook(module, grad_input, grad_output):
        gradcam_data["grad"] = grad_output[0].detach()

    target_layer = model.swinv2.encoder.layers[-1].blocks[-1].layernorm_after
    fwd_handle = target_layer.register_forward_hook(fwd_hook)
    bwd_handle = target_layer.register_full_backward_hook(bwd_hook)

    local_files = []
    for idx in selected:
        try:
            item = test_ds[idx]
            img = item["image"].convert("RGB")
            label = item["label"]
            inputs = processor(img, return_tensors="pt")
            img_tensor = inputs["pixel_values"].to(device).requires_grad_(True)
            model.zero_grad()
            outputs = model(pixel_values=img_tensor)
            target_class = int(y_pred[idx])
            score = outputs.logits[0, target_class]
            score.backward()

            feat = gradcam_data["feat"][0]      # [H*W, C]
            grads = gradcam_data["grad"][0]   # [H*W, C]
            weights = grads.mean(dim=0, keepdim=True)
            cam = torch.matmul(feat, weights.t()).squeeze()  # [H*W]
            H = W = int(math.sqrt(cam.shape[0]))
            cam = cam.reshape(H, W)
            cam = F.relu(cam)
            cam = cam - cam.min()
            cam = cam / (cam.max() + 1e-8)
            cam = F.interpolate(cam.unsqueeze(0).unsqueeze(0), size=(256, 256), mode="bilinear", align_corners=False)
            cam = cam.squeeze().cpu().numpy()

            img_np = img_tensor.squeeze().detach().cpu().permute(1,2,0).numpy()
            img_np = (img_np - img_np.min()) / (img_np.max() - img_np.min() + 1e-8)

            plt.figure(figsize=(6,6))
            plt.imshow(img_np)
            plt.imshow(cam, cmap="jet", alpha=0.5)
            pred_name = id2label.get(target_class, str(target_class))
            true_name = id2label.get(label, str(label))
            status = "CORRECT" if y_true[idx] == y_pred[idx] else "WRONG"
            plt.title(f"{status}: Pred={pred_name} | True={true_name}")
            plt.axis("off")
            fname = f"gradcam_{status}_sample{idx}_{pred_name}_vs_{true_name}.png"
            fpath = os.path.join(OUTPUT_DIR, fname)
            plt.savefig(fpath, bbox_inches="tight", dpi=150)
            plt.close()
            local_files.append(fpath)
            print(f"  Saved: {fpath}")
        except Exception as e:
            print(f"  Skipped sample {idx}: {e}")
            traceback.print_exc()

    fwd_handle.remove()
    bwd_handle.remove()

    print(f"\n{'='*60}")
    print(f"Generated {len(local_files)} Grad-CAM images in {OUTPUT_DIR}/")
    print(f"{'='*60}")

    # Upload to Hub
    print("\nUploading to Hugging Face Hub...")
    api = HfApi()
    uploaded = 0
    for fpath in local_files:
        fname = os.path.basename(fpath)
        try:
            api.upload_file(
                path_or_fileobj=fpath,
                path_in_file=f"gradcam/{fname}",
                repo_id=REPO_ID,
                repo_type="model"
            )
            print(f"  Uploaded: gradcam/{fname}")
            uploaded += 1
        except Exception as e:
            print(f"  Failed to upload {fname}: {e}")

    print(f"\nDone! Uploaded {uploaded}/{len(local_files)} images to https://huggingface.co/{REPO_ID}/tree/main/gradcam/")

if __name__ == "__main__":
    main()