| """Side-by-side HiResCAM comparison across the 3 training regimes.
|
|
|
| Picks 4 representative images, runs HiResCAM with each of:
|
| - best_original.pth (raw baseline)
|
| - best_sam.pth (SAM bg-removal)
|
| - best_strongaug.pth (strong augmentation)
|
|
|
| Builds a 4x4 grid (rows = images, cols = original | 3 CAMs).
|
| """
|
| import sys, os
|
| sys.path.insert(0, "/mnt/d/SpiceNet" if os.path.exists("/mnt/d/SpiceNet") else "D:/SpiceNet")
|
|
|
| import numpy as np
|
| import torch
|
| import cv2
|
| import matplotlib
|
| matplotlib.use("Agg")
|
| import matplotlib.pyplot as plt
|
| from PIL import Image
|
| import albumentations as A
|
| from albumentations.pytorch import ToTensorV2
|
|
|
| import config
|
| from src.model import load_checkpoint
|
| from src.gradcam import GradCAM
|
| from src.features import extract_all as _extract_all
|
|
|
|
|
| _RESIZE = int(config.IMG_SIZE * 256 / 224)
|
| _CROP = config.IMG_SIZE
|
|
|
|
|
| def model_view(img_path):
|
| img = np.array(Image.open(img_path).convert("RGB"))
|
| crop_tf = A.Compose([A.Resize(_RESIZE, _RESIZE), A.CenterCrop(_CROP, _CROP)])
|
| display = crop_tf(image=img)["image"]
|
| norm_tf = A.Compose([A.Normalize(mean=config.IMG_MEAN, std=config.IMG_STD), ToTensorV2()])
|
| tensor = norm_tf(image=display)["image"].unsqueeze(0)
|
| tex_np, col_np = _extract_all(img)
|
| tex = torch.from_numpy(tex_np).unsqueeze(0).float()
|
| col = torch.from_numpy(col_np).unsqueeze(0).float()
|
| return tensor, display, tex, col
|
|
|
|
|
| def compute_cam(model, tensor, tex, col, device, smooth_sigma=1.2):
|
| gcam = GradCAM(model, mode="fusion", algorithm="hirescam",
|
| smooth_sigma=smooth_sigma)
|
| tensor = tensor.to(device); tex = tex.to(device); col = col.to(device)
|
| cam, pred = gcam(tensor, tex=tex, col=col)
|
| gcam.remove()
|
| return cam, pred
|
|
|
|
|
| def overlay(display, cam, smooth_sigma):
|
| heat = cv2.resize(cam, (display.shape[1], display.shape[0]),
|
| interpolation=cv2.INTER_CUBIC)
|
| if smooth_sigma > 0:
|
| k = max(3, int(2 * round(3 * smooth_sigma) + 1))
|
| heat = cv2.GaussianBlur(heat, (k, k), smooth_sigma)
|
| heat = heat - heat.min()
|
| if heat.max() > 0:
|
| heat = heat / heat.max()
|
| heat_u8 = cv2.applyColorMap(np.uint8(255 * heat), cv2.COLORMAP_JET)
|
| heat_u8 = cv2.cvtColor(heat_u8, cv2.COLOR_BGR2RGB)
|
| return cv2.addWeighted(display, 0.5, heat_u8, 0.5, 0)
|
|
|
|
|
| def main():
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
|
| samples = [
|
| ("Spice_Spectrum/turmeric/turmeric_10.jpg", "turmeric"),
|
| ("Spice_Spectrum/cumin/cumin_50.jpg", "cumin"),
|
| ("Spice_Spectrum/cinnamon/cinnamon_100.jpg", "cinnamon"),
|
| ("Spice_Spectrum/saffron/saffron_200.jpg", "saffron"),
|
| ]
|
|
|
| pipelines = [
|
| ("Raw baseline", "best_original.pth"),
|
| ("SAM bg-removal", "best_sam.pth"),
|
| ("Strong aug", "best_strongaug.pth"),
|
| ]
|
|
|
| fig, axes = plt.subplots(len(samples), 1 + len(pipelines),
|
| figsize=(3 * (1 + len(pipelines)), 3 * len(samples)))
|
|
|
| for row, (path, lbl) in enumerate(samples):
|
| tensor, display, tex, col = model_view(path)
|
| axes[row, 0].imshow(display)
|
| axes[row, 0].set_title("Input" if row == 0 else "", fontsize=10)
|
| axes[row, 0].set_ylabel(lbl, fontsize=10)
|
| axes[row, 0].set_xticks([]); axes[row, 0].set_yticks([])
|
|
|
| for col_idx, (pname, ckpt_name) in enumerate(pipelines, start=1):
|
| ckpt_path = config.CHECKPOINT_DIR / ckpt_name
|
| if not ckpt_path.exists():
|
| axes[row, col_idx].text(0.5, 0.5, f"missing: {ckpt_name}",
|
| ha="center", va="center")
|
| axes[row, col_idx].set_xticks([]); axes[row, col_idx].set_yticks([])
|
| continue
|
| model, *_ = load_checkpoint(str(ckpt_path), device)
|
| cam, pred_idx = compute_cam(model, tensor, tex, col, device)
|
| ov = overlay(display, cam, smooth_sigma=1.2)
|
| axes[row, col_idx].imshow(ov)
|
| pred_name = config.CLASSES[pred_idx]
|
| color = "green" if pred_name == lbl else "red"
|
| axes[row, col_idx].set_title(
|
| f"{pname}\npred: {pred_name}" if row == 0 else f"pred: {pred_name}",
|
| fontsize=9, color=color,
|
| )
|
| axes[row, col_idx].set_xticks([]); axes[row, col_idx].set_yticks([])
|
| del model
|
| if device.type == "cuda":
|
| torch.cuda.empty_cache()
|
|
|
| plt.suptitle("HiResCAM comparison — same images, 3 training regimes", fontsize=13)
|
| plt.tight_layout()
|
| out = config.OUTPUT_DIR / "cam_comparison_3regimes.png"
|
| plt.savefig(out, dpi=150, bbox_inches="tight")
|
| plt.close()
|
| print(f"Saved -> {out}")
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
