evaluate_and_gradcam.py

"""
Thyroid Ultrasound Evaluation + Grad-CAM Visualization
Evaluates model on test set and generates attention visualizations.
"""
import os, sys, io, math, json, random, warnings, base64, 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,
    Trainer, TrainingArguments, DefaultDataCollator
)
from sklearn.metrics import (
    accuracy_score, precision_recall_fscore_support, roc_auc_score, confusion_matrix
)

os.environ["TRACKIO_SPACE_ID"] = ""
os.environ["TRACKIO_PROJECT"] = ""

HF_USERNAME = "Johnyquest7"
DATASET_NAME = "BTX24/thyroid-cancer-classification-ultrasound-dataset"
MODEL_NAME = f"{HF_USERNAME}/ML-Inter_thyroid"
OUTPUT_DIR = "./eval_outputs"
SEED = 42
MAX_SAMPLES_GRADCAM = 20

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

def main():
    print("=" * 60)
    print("Thyroid Ultrasound Model Evaluation + Grad-CAM")
    print("=" * 60)
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"\nDevice: {device}")
    print(f"Loading model: {MODEL_NAME}")

    try:
        processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
        model = AutoModelForImageClassification.from_pretrained(MODEL_NAME).to(device).eval()
    except Exception as e:
        print(f"Model loading failed: {e}")
        sys.exit(1)

    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)}")

    id2label = model.config.id2label
    label2id = model.config.label2id

    def transform(examples):
        images = [img.convert("RGB") if img.mode != "RGB" else img for img in examples["image"]]
        return processor(images, return_tensors="pt")

    test_ds.set_transform(transform)

    # Evaluate
    print("\nRunning evaluation...")
    args = TrainingArguments(
        output_dir="/tmp/eval", per_device_eval_batch_size=16,
        remove_unused_columns=False, disable_tqdm=True,
        logging_strategy="steps", logging_first_step=True,
        report_to=[]
    )
    trainer = Trainer(model=model, args=args, data_collator=DefaultDataCollator(),
                      eval_dataset=test_ds)
    metrics = trainer.evaluate()
    print(f"\nRaw metrics: {metrics}")

    # Collect predictions
    all_logits, all_labels = [], []
    for i in range(0, len(test_ds), 16):
        batch = test_ds[i:i+16]
        inputs = {k: torch.stack([v for v in batch[k]]).to(device) if isinstance(batch[k][0], torch.Tensor) else None
                  for k in batch if k in processor.model_input_names or k == "pixel_values"}
        if "pixel_values" in inputs and inputs["pixel_values"] is not None:
            with torch.no_grad():
                outputs = model(pixel_values=inputs["pixel_values"])
            all_logits.extend(outputs.logits.cpu().numpy())
            all_labels.extend(batch["label"])

    y_true = np.array(all_labels)
    y_logits = np.array(all_logits)
    y_pred = np.argmax(y_logits, axis=1)
    probs = F.softmax(torch.from_numpy(y_logits), dim=1).numpy()

    acc = accuracy_score(y_true, y_pred)
    prec, rec, f1, _ = precision_recall_fscore_support(y_true, y_pred, average="weighted")
    try:
        auc = roc_auc_score(y_true, probs[:, 1])
    except:
        auc = roc_auc_score(y_true, probs[:, 0])
    cm = confusion_matrix(y_true, y_pred)

    final = {
        "test_accuracy": float(acc),
        "test_weighted_f1": float(f1),
        "test_weighted_precision": float(prec),
        "test_weighted_recall": float(rec),
        "test_roc_auc": float(auc),
        "test_confusion_matrix": cm.tolist(),
        "eval_loss": float(metrics.get("eval_loss", 0)),
    }
    print(f"\n{'='*60}")
    print("FINAL TEST METRICS")
    print(f"{'='*60}")
    for k, v in final.items():
        print(f"  {k}: {v}")
    json.dump(final, open(f"{OUTPUT_DIR}/test_metrics.json", "w"), indent=2)
    print(f"\nSaved to {OUTPUT_DIR}/test_metrics.json")

    # Grad-CAM: collect misclassified and correct
    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[:min(5, len(correct_idx))] + incorrect_idx[:min(5, len(incorrect_idx))]
    print(f"\nGenerating Grad-CAM for {len(selected)} samples ({len(correct_idx[:5])} correct, {len(incorrect_idx[:5])} incorrect)...")

    # Hook into last stage norm of Swin
    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)

    for idx in selected[:MAX_SAMPLES_GRADCAM]:
        try:
            sample = test_ds[idx]
            label = sample["label"]
            img_tensor = sample["pixel_values"].unsqueeze(0).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]
            grads = gradcam_data["grad"][0]
            if feat.dim() == 3:  # Swin output (H*W, C)
                weights = grads.mean(dim=0, keepdim=True)
                cam = torch.matmul(feat, weights.t()).squeeze()
                H = W = int(math.sqrt(cam.shape[0]))
                cam = cam.reshape(H, W)
            else:
                weights = grads.mean(dim=(0,1), keepdim=True)
                cam = (feat * weights).sum(dim=-1).squeeze()

            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()

            # Overlay
            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)
            plt.title(f"Pred: {id2label[target_class]} | True: {id2label[label]}")
            plt.axis("off")
            fname = f"{OUTPUT_DIR}/gradcam_sample_{idx}_pred{id2label[target_class]}_true{id2label[label]}.png"
            plt.savefig(fname, bbox_inches="tight", dpi=150)
            plt.close()
            print(f"  Saved {fname}")
        except Exception as e:
            print(f"  Skipped sample {idx}: {e}")
            traceback.print_exc()

    fwd_handle.remove()
    bwd_handle.remove()

    # Push outputs to Hub as a dataset or files
    print("\nEvaluation complete.")
    print(f"Results saved to {OUTPUT_DIR}/")

if __name__ == "__main__":
    main()