evaluate.py
4,030 bytes
| 1 | import argparse |
|---|---|
| 2 | import json |
| 3 | from pathlib import Path |
| 4 | |
| 5 | import matplotlib |
| 6 | matplotlib.use("Agg") |
| 7 | import matplotlib.pyplot as plt |
| 8 | import numpy as np |
| 9 | import pandas as pd |
| 10 | import timm |
| 11 | import torch |
| 12 | from sklearn.metrics import confusion_matrix, f1_score |
| 13 | from torch.utils.data import DataLoader |
| 14 | |
| 15 | from .data import CountryDataset, build_transforms |
| 16 | from .model import load_checkpoint |
| 17 | |
| 18 | |
| 19 | @torch.no_grad() |
| 20 | def collect_predictions(model, loader, device): |
| 21 | preds, tops, targets = [], [], [] |
| 22 | for images, labels in loader: |
| 23 | images = images.to(device, non_blocking=True) |
| 24 | logits = model(images) |
| 25 | k = min(5, logits.size(1)) |
| 26 | preds.append(logits.argmax(1).cpu()) |
| 27 | tops.append(logits.topk(k, dim=1).indices.cpu()) |
| 28 | targets.append(labels) |
| 29 | return torch.cat(preds).numpy(), torch.cat(tops).numpy(), torch.cat(targets).numpy() |
| 30 | |
| 31 | |
| 32 | def main(): |
| 33 | ap = argparse.ArgumentParser(description="Hindab mudelit testihulgal") |
| 34 | ap.add_argument("--checkpoint", default="runs/clip_vit_b16/best.pt") |
| 35 | ap.add_argument("--split-csv", default=None, help="vaikimisi test.csv checkpointi kaustast") |
| 36 | ap.add_argument("--batch-size", type=int, default=64) |
| 37 | ap.add_argument("--num-workers", type=int, default=2) |
| 38 | ap.add_argument("--plot-top", type=int, default=30, help="mitu suurima toega riiki joonisele") |
| 39 | args = ap.parse_args() |
| 40 | |
| 41 | ckpt_dir = Path(args.checkpoint).parent |
| 42 | split_csv = Path(args.split_csv) if args.split_csv else ckpt_dir / "test.csv" |
| 43 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| 44 | |
| 45 | model, classes = load_checkpoint(args.checkpoint, device) |
| 46 | class_to_idx = {c: i for i, c in enumerate(classes)} |
| 47 | |
| 48 | df = pd.read_csv(split_csv) |
| 49 | df = df[df["country"].isin(class_to_idx)].reset_index(drop=True) |
| 50 | data_cfg = timm.data.resolve_model_data_config(model.backbone) |
| 51 | ds = CountryDataset( |
| 52 | df, class_to_idx, |
| 53 | build_transforms(data_cfg["mean"], data_cfg["std"], data_cfg["input_size"][-1], train=False), |
| 54 | ) |
| 55 | loader = DataLoader(ds, batch_size=args.batch_size, num_workers=args.num_workers) |
| 56 | |
| 57 | preds, tops, targets = collect_predictions(model, loader, device) |
| 58 | top1 = float((preds == targets).mean()) |
| 59 | top5 = float((tops == targets[:, None]).any(1).mean()) |
| 60 | macro_f1 = float(f1_score(targets, preds, average="macro")) |
| 61 | |
| 62 | metrics = {"top1": top1, "top5": top5, "macro_f1": macro_f1, |
| 63 | "n_test": int(len(targets)), "n_classes": len(classes)} |
| 64 | with open(ckpt_dir / "metrics.json", "w", encoding="utf-8") as f: |
| 65 | json.dump(metrics, f, indent=2) |
| 66 | print(json.dumps(metrics, indent=2)) |
| 67 | |
| 68 | per_country = ( |
| 69 | pd.DataFrame({"country": [classes[i] for i in targets], "correct": preds == targets}) |
| 70 | .groupby("country") |
| 71 | .agg(accuracy=("correct", "mean"), support=("correct", "size")) |
| 72 | .sort_values("accuracy") |
| 73 | ) |
| 74 | per_country.to_csv(ckpt_dir / "per_country.csv") |
| 75 | print("Nõrgimad riigid:") |
| 76 | print(per_country.head(10).to_string()) |
| 77 | |
| 78 | cm = confusion_matrix(targets, preds, labels=range(len(classes))) |
| 79 | pd.DataFrame(cm, index=classes, columns=classes).to_csv(ckpt_dir / "confusion_matrix.csv") |
| 80 | |
| 81 | # Joonisele ainult suurima toega riigid, muidu on maatriks loetamatu. |
| 82 | top_names = per_country.sort_values("support", ascending=False).head(args.plot_top).index |
| 83 | idx = [class_to_idx[c] for c in top_names] |
| 84 | sub = cm[np.ix_(idx, idx)].astype(np.float64) |
| 85 | sub = sub / sub.sum(axis=1, keepdims=True).clip(min=1) |
| 86 | fig, ax = plt.subplots(figsize=(12, 10)) |
| 87 | ax.imshow(sub, cmap="Blues", vmin=0, vmax=1) |
| 88 | ax.set_xticks(range(len(idx)), top_names, rotation=90, fontsize=7) |
| 89 | ax.set_yticks(range(len(idx)), top_names, fontsize=7) |
| 90 | ax.set_xlabel("Ennustatud") |
| 91 | ax.set_ylabel("Tegelik") |
| 92 | ax.set_title(f"Confusion matrix, {args.plot_top} suurima toega riiki (reanormeeritud)") |
| 93 | fig.tight_layout() |
| 94 | fig.savefig(ckpt_dir / "confusion_matrix.png", dpi=150) |
| 95 | print(f"Tulemused kaustas {ckpt_dir}") |
| 96 | |
| 97 | |
| 98 | if __name__ == "__main__": |
| 99 | main() |
| 100 | |