File size: 4,829 Bytes
dda557a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 | """Evaluate the trained model on the test set and dump metrics.json.
Run after `train.py` has produced `model/saved/brain_tumor_model.pth`.
Outputs:
model/saved/metrics.json - test accuracy, per-class precision/recall/F1, confusion matrix
model/saved/confusion_matrix.png - heatmap visualization (if matplotlib available)
"""
from __future__ import annotations
import json
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
from sklearn.metrics import (
classification_report,
confusion_matrix,
)
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import ImageFolder
from architecture import CLASS_NAMES, IMG_SIZE, build_model
PROJECT_ROOT = Path(__file__).resolve().parent.parent
TEST_DIR = PROJECT_ROOT / "data" / "raw" / "Testing"
SAVE_DIR = PROJECT_ROOT / "model" / "saved"
CHECKPOINT = SAVE_DIR / "brain_tumor_model.pth"
METRICS_PATH = SAVE_DIR / "metrics.json"
CM_PLOT_PATH = SAVE_DIR / "confusion_matrix.png"
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]
def _save_confusion_plot(cm: np.ndarray, labels: list[str]) -> None:
try:
import matplotlib.pyplot as plt
except ImportError:
return
fig, ax = plt.subplots(figsize=(6, 5))
im = ax.imshow(cm, cmap="Blues")
ax.set_xticks(range(len(labels)))
ax.set_yticks(range(len(labels)))
ax.set_xticklabels(labels, rotation=45, ha="right")
ax.set_yticklabels(labels)
ax.set_xlabel("Predicted")
ax.set_ylabel("True")
ax.set_title("Confusion Matrix (test set)")
for i in range(cm.shape[0]):
for j in range(cm.shape[1]):
ax.text(j, i, str(cm[i, j]), ha="center", va="center", color="black")
fig.colorbar(im, ax=ax)
fig.tight_layout()
fig.savefig(CM_PLOT_PATH, dpi=150)
plt.close(fig)
print(f"Confusion matrix plot saved to {CM_PLOT_PATH}")
def main() -> None:
if not CHECKPOINT.exists():
raise SystemExit(f"Checkpoint not found at {CHECKPOINT}. Run train.py first.")
# Match the eval transform used during training: Resize(IMG_SIZE+24) → CenterCrop(IMG_SIZE).
eval_transform = transforms.Compose(
[
transforms.Resize((IMG_SIZE + 24, IMG_SIZE + 24)),
transforms.CenterCrop(IMG_SIZE),
transforms.ToTensor(),
transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD),
]
)
test_ds = ImageFolder(TEST_DIR, transform=eval_transform)
if test_ds.classes != CLASS_NAMES:
raise SystemExit(
f"Class order mismatch: {test_ds.classes} vs expected {CLASS_NAMES}"
)
loader = DataLoader(
test_ds, batch_size=32, shuffle=False, num_workers=4, pin_memory=True
)
print(f"Device: {DEVICE} | test set: {len(test_ds)} images")
model = build_model(pretrained=False).to(DEVICE)
ckpt = torch.load(CHECKPOINT, map_location=DEVICE)
model.load_state_dict(ckpt["model_state_dict"])
model.eval()
all_preds: list[int] = []
all_targets: list[int] = []
# Match training criterion so the reported loss is comparable to train_v2.log.
criterion = nn.CrossEntropyLoss(label_smoothing=0.05)
loss_sum, total = 0.0, 0
with torch.no_grad():
for inputs, targets in loader:
inputs = inputs.to(DEVICE, non_blocking=True)
targets = targets.to(DEVICE, non_blocking=True)
logits = model(inputs)
loss_sum += criterion(logits, targets).item() * inputs.size(0)
total += inputs.size(0)
preds = logits.argmax(dim=1).cpu().numpy()
all_preds.extend(preds.tolist())
all_targets.extend(targets.cpu().numpy().tolist())
acc = float(np.mean(np.array(all_preds) == np.array(all_targets)))
test_loss = loss_sum / total
cm = confusion_matrix(all_targets, all_preds).tolist()
report = classification_report(
all_targets,
all_preds,
target_names=CLASS_NAMES,
digits=4,
output_dict=True,
)
metrics = dict(
test_accuracy=acc,
test_loss=test_loss,
num_samples=total,
class_names=CLASS_NAMES,
confusion_matrix=cm,
classification_report=report,
checkpoint=str(CHECKPOINT.relative_to(PROJECT_ROOT)),
)
METRICS_PATH.write_text(json.dumps(metrics, indent=2))
print(f"Test accuracy: {acc*100:.2f}% | loss: {test_loss:.4f}")
print(f"Per-class F1: " + ", ".join(
f"{c}={report[c]['f1-score']:.3f}" for c in CLASS_NAMES
))
print(f"Confusion matrix:\n{np.array(cm)}")
print(f"Metrics written to {METRICS_PATH}")
_save_confusion_plot(np.array(cm), CLASS_NAMES)
if __name__ == "__main__":
main()
|