Update test_new.py

test_new.py

@@ -1,319 +1,319 @@
-import argparse
-import numpy as np
-import torch
-import torch.backends.cudnn as cudnn
-import os
-import warnings
-import json
-from pathlib import Path
-
-from timm.models import create_model
-
-import my_models  # registra TALL_SWIN
-import utils
-
-from video_dataset import VideoDataSet
-from video_dataset_aug import get_augmentor, build_dataflow
-from video_dataset_config import get_dataset_config, DATASET_CONFIG
-
-from sklearn.metrics import (
-    accuracy_score, balanced_accuracy_score,
-    precision_recall_fscore_support,
-    confusion_matrix, classification_report,
-    roc_auc_score, roc_curve,
-    average_precision_score, precision_recall_curve
-)
-import matplotlib.pyplot as plt
-
-warnings.filterwarnings("ignore", category=UserWarning)
-
-
-def get_args_parser():
-    parser = argparse.ArgumentParser('DeiT evaluation script', add_help=False)
-
-    parser.add_argument('--model', default='TALL_SWIN', type=str)
-    parser.add_argument('--model_name', default="TALL_SWIN")
-    parser.add_argument('--batch-size', default=2, type=int)
-
-    # Dataset parameters
-    parser.add_argument('--data_txt_dir', type=str, default='##path_for_dataset_txt##')
-    parser.add_argument('--data_dir', type=str, default="##path_for_dataset##")
-    parser.add_argument('--dataset', default='ffpp', choices=list(DATASET_CONFIG.keys()))
-    parser.add_argument('--duration', default=1, type=int)
-    parser.add_argument('--frames_per_group', default=1, type=int)
-    parser.add_argument('--threed_data', default=False)
-    parser.add_argument('--input_size', default=224, type=int)
-    parser.add_argument('--disable_scaleup', action='store_true')
-    parser.add_argument('--random_sampling', action='store_true')
-    parser.add_argument('--dense_sampling', default=True)
-    parser.add_argument('--augmentor_ver', default='v1', type=str, choices=['v1', 'v2'])
-    parser.add_argument('--scale_range', default=[256, 320], type=int, nargs="+")
-    parser.add_argument('--modality', default='rgb', type=str)
-    parser.add_argument('--use_lmdb', default=False)
-    parser.add_argument('--use_pyav', default=False)
-
-    # temporal module / model params
-    parser.add_argument('--pretrained', action='store_true', default=False)
-    parser.add_argument('--temporal_module_name', default=None, type=str,
-                        choices=['ResNet3d', 'TAM', 'TTAM', 'TSM', 'TTSM', 'MSA'])
-    parser.add_argument('--temporal_attention_only', action='store_true', default=False)
-    parser.add_argument('--no_token_mask', action='store_true', default=False)
-    parser.add_argument('--temporal_heads_scale', default=1.0, type=float)
-    parser.add_argument('--temporal_mlp_scale', default=1.0, type=float)
-    parser.add_argument('--rel_pos', action='store_true', default=False)
-    parser.add_argument('--temporal_pooling', type=str, default=None,
-                        choices=['avg', 'max', 'conv', 'depthconv'])
-    parser.add_argument('--bottleneck', default=None, choices=['regular', 'dw'])
-
-    parser.add_argument('--window_size', default=7, type=int)
-    parser.add_argument('--thumbnail_rows', default=3, type=int)
-    parser.add_argument('--hpe_to_token', default=False, action='store_true')
-
-    parser.add_argument('--drop', type=float, default=0.0)
-    parser.add_argument('--drop-path', type=float, default=0.1)
-    parser.add_argument('--drop-block', type=float, default=None)
-
-    # runtime
-    parser.add_argument('--output_dir', default="./output")
-    parser.add_argument('--device', default='cuda')
-    parser.add_argument('--seed', default=42, type=int)
-    parser.add_argument('--num_workers', default=8, type=int)
-
-    parser.add_argument('--num_crops', default=1, type=int, choices=[1, 3, 5, 10])
-    parser.add_argument('--num_clips', default=3, type=int)
-
-    parser.add_argument('--world_size', default=1, type=int)
-    parser.add_argument("--local_rank", type=int)
-    parser.add_argument('--dist_url', default='env://')
-
-    # checkpoint
-    parser.add_argument('--initial_checkpoint', type=str, default='',
-                        help='path do .pth/.pth.tar com checkpoint (espera key "model")')
-
-    parser.add_argument('--threshold', type=float, default=0.5,
-                        help='threshold para decidir classe 1 (fake) a partir de prob[:,1]')
-    parser.add_argument('--metrics_out', default='', type=str,
-                        help='pasta pra salvar metrics.json e plots (default: output_dir)')
-    parser.add_argument('--save_plots', action='store_true',
-                        help='salvar cm.png / roc.png / pr.png')
-
-    return parser
-
-
-@torch.no_grad()
-def eval_with_outputs(data_loader, model, device, threshold: float = 0.5):
-    model.eval()
-    y_true, y_score, y_pred = [], [], []
-
-    thr = float(threshold)
-
-    for samples, targets in data_loader:
-        samples = samples.to(device, non_blocking=True)
-        targets = targets.to(device, non_blocking=True)
-
-        logits = model(samples)  # [B,2] ou [B*K,2]
-
-        # se logits veio por-clip, agrega por vídeo
-        B = targets.shape[0]
-        if logits.shape[0] != B:
-            if logits.shape[0] % B != 0:
-                raise RuntimeError(
-                    f"logits batch ({logits.shape[0]}) não é múltiplo do target batch ({B})."
-                )
-            K = logits.shape[0] // B
-            logits = logits.view(B, K, -1).mean(dim=1)  # [B,2]
-
-        probs = torch.softmax(logits, dim=1)  # [B,2]
-        p1 = probs[:, 1]                      # score da classe 1 (fake)
-
-        # >>> AQUI é o THRESHOLD <<<
-        hat = (p1 >= thr).long()
-
-        y_true.append(targets.detach().cpu().numpy())
-        y_score.append(p1.detach().cpu().numpy())
-        y_pred.append(hat.detach().cpu().numpy())
-
-    y_true = np.concatenate(y_true).astype(int)
-    y_score = np.concatenate(y_score).astype(float)
-    y_pred = np.concatenate(y_pred).astype(int)
-    return y_true, y_score, y_pred
-
-
-def plot_confusion(cm, out_path):
-    plt.figure(figsize=(6, 5))
-    plt.imshow(cm)
-    plt.title("Confusion Matrix")
-    plt.xlabel("Predicted")
-    plt.ylabel("True")
-    for (i, j), v in np.ndenumerate(cm):
-        plt.text(j, i, str(v), ha="center", va="center")
-    plt.tight_layout()
-    plt.savefig(out_path, dpi=200)
-    plt.close()
-
-
-def plot_roc(y, scores, out_path):
-    fpr, tpr, _ = roc_curve(y, scores)
-    auc = roc_auc_score(y, scores)
-    plt.figure(figsize=(7, 6))
-    plt.plot(fpr, tpr, label=f"AUC={auc:.4f}")
-    plt.plot([0, 1], [0, 1], "--", label="Chance")
-    plt.xlabel("FPR")
-    plt.ylabel("TPR")
-    plt.legend(loc="best")
-    plt.tight_layout()
-    plt.savefig(out_path, dpi=200)
-    plt.close()
-
-
-def plot_pr(y, scores, out_path):
-    p, r, _ = precision_recall_curve(y, scores)
-    ap = average_precision_score(y, scores)
-    plt.figure(figsize=(7, 6))
-    plt.plot(r, p, label=f"AP={ap:.4f}")
-    plt.xlabel("Recall")
-    plt.ylabel("Precision")
-    plt.legend(loc="best")
-    plt.tight_layout()
-    plt.savefig(out_path, dpi=200)
-    plt.close()
-
-
-def main(args):
-    utils.init_distributed_mode(args)
-    print(args)
-
-    device = torch.device(args.device)
-
-    seed = args.seed + utils.get_rank()
-    torch.manual_seed(seed)
-    np.random.seed(seed)
-    cudnn.benchmark = True
-
-    num_classes, train_list_name, val_list_name, test_list_name, filename_seperator, image_tmpl, filter_video, label_file = \
-        get_dataset_config(args.dataset, args.use_lmdb)
-
-    args.num_classes = num_classes
-    args.input_channels = 3 if args.modality == 'rgb' else 2 * 5
-
-    print(f"Creating model: {args.model}")
-    model = create_model(
-        args.model,
-        pretrained=args.pretrained,
-        duration=args.duration,
-        hpe_to_token=args.hpe_to_token,
-        rel_pos=args.rel_pos,
-        window_size=args.window_size,
-        thumbnail_rows=args.thumbnail_rows,
-        token_mask=not args.no_token_mask,
-        online_learning=False,
-        num_classes=args.num_classes,
-        drop_rate=args.drop,
-        drop_path_rate=args.drop_path,
-        drop_block_rate=args.drop_block,
-        use_checkpoint=False
-    )
-    model.to(device)
-
-    # mean/std
-    if args.distributed:
-        mean = (0.5, 0.5, 0.5) if 'mean' not in model.module.default_cfg else model.module.default_cfg['mean']
-        std = (0.5, 0.5, 0.5) if 'std' not in model.module.default_cfg else model.module.default_cfg['std']
-    else:
-        mean = (0.5, 0.5, 0.5) if 'mean' not in model.default_cfg else model.default_cfg['mean']
-        std = (0.5, 0.5, 0.5) if 'std' not in model.default_cfg else model.default_cfg['std']
-
-    # dataset (val list)
-    video_data_cls = VideoDataSet
-    val_list = os.path.join(args.data_txt_dir, val_list_name)
-
-    val_augmentor = get_augmentor(
-        False, args.input_size, mean, std, args.disable_scaleup,
-        threed_data=args.threed_data, version=args.augmentor_ver,
-        scale_range=args.scale_range, num_clips=args.num_clips,
-        num_crops=args.num_crops, dataset=args.dataset
-    )
-
-    dataset_val = video_data_cls(
-        args.data_dir, val_list,
-        args.duration, args.frames_per_group,
-        num_clips=args.num_clips,
-        modality=args.modality,
-        dense_sampling=args.dense_sampling,
-        image_tmpl=image_tmpl,
-        transform=val_augmentor,
-        is_train=False, test_mode=False,
-        seperator=filename_seperator, filter_video=filter_video
-    )
-
-    data_loader_val = build_dataflow(
-        dataset_val, is_train=False, batch_size=args.batch_size,
-        workers=args.num_workers, is_distributed=args.distributed
-    )
-
-    if not args.initial_checkpoint:
-        raise RuntimeError("Passe --initial_checkpoint apontando pro checkpoint do modelo.")
-
-    checkpoint = torch.load(args.initial_checkpoint, map_location='cpu')
-    # muitos checkpoints vêm como {"model": state_dict, ...}
-    if isinstance(checkpoint, dict) and "model" in checkpoint:
-        utils.load_checkpoint(model, checkpoint["model"])
-    else:
-        # se for state_dict direto
-        model.load_state_dict(checkpoint, strict=False)
-
-    # eval
-    y_true, y_score, y_pred = eval_with_outputs(
-        data_loader_val, model, device, threshold=args.threshold
-    )
-
-    acc = accuracy_score(y_true, y_pred)
-    bacc = balanced_accuracy_score(y_true, y_pred)
-    prec, rec, f1, _ = precision_recall_fscore_support(
-        y_true, y_pred, average="binary", zero_division=0
-    )
-    cm = confusion_matrix(y_true, y_pred)
-
-    roc_auc = roc_auc_score(y_true, y_score)
-    pr_auc = average_precision_score(y_true, y_score)
-
-    print(f"\nN={len(y_true)} | thr={args.threshold:.3f}")
-    print(f"acc={acc:.4f} | bacc={bacc:.4f} | prec={prec:.4f} | rec={rec:.4f} | f1={f1:.4f} | roc_auc={roc_auc:.4f} | pr_auc={pr_auc:.4f}")
-    print(classification_report(y_true, y_pred, digits=4, zero_division=0))
-
-    outdir = args.metrics_out.strip() if args.metrics_out else args.output_dir
-    os.makedirs(outdir, exist_ok=True)
-
-    out_json = {
-        "threshold": float(args.threshold),
-        "acc": float(acc),
-        "balanced_acc": float(bacc),
-        "precision": float(prec),
-        "recall": float(rec),
-        "f1": float(f1),
-        "roc_auc": float(roc_auc),
-        "pr_auc": float(pr_auc),
-        "confusion_matrix": cm.tolist(),
-        "n": int(len(y_true)),
-    }
-    with open(os.path.join(outdir, "metrics.json"), "w", encoding="utf-8") as f:
-        json.dump(out_json, f, indent=2)
-
-    np.savez(os.path.join(outdir, "eval_outputs.npz"),
-             y_true=y_true, y_score=y_score, y_pred=y_pred)
-
-    if args.save_plots:
-        plot_confusion(cm, os.path.join(outdir, "cm.png"))
-        plot_roc(y_true, y_score, os.path.join(outdir, "roc.png"))
-        plot_pr(y_true, y_score, os.path.join(outdir, "pr.png"))
-        print(f"\n✔ Plots + metrics saved in: {os.path.abspath(outdir)}")
-    else:
-        print(f"\n✔ Metrics saved in: {os.path.abspath(os.path.join(outdir, 'metrics.json'))}")
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser('DeiT evaluation script', parents=[get_args_parser()])
-    args = parser.parse_args()
-    if args.output_dir:
-        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
-    main(args)
+import argparse
+import numpy as np
+import torch
+import torch.backends.cudnn as cudnn
+import os
+import warnings
+import json
+from pathlib import Path
+
+from timm.models import create_model
+
+import my_models  # registers TALL_SWIN
+import utils
+
+from video_dataset import VideoDataSet
+from video_dataset_aug import get_augmentor, build_dataflow
+from video_dataset_config import get_dataset_config, DATASET_CONFIG
+
+from sklearn.metrics import (
+    accuracy_score, balanced_accuracy_score,
+    precision_recall_fscore_support,
+    confusion_matrix, classification_report,
+    roc_auc_score, roc_curve,
+    average_precision_score, precision_recall_curve
+)
+import matplotlib.pyplot as plt
+
+warnings.filterwarnings("ignore", category=UserWarning)
+
+
+def get_args_parser():
+    parser = argparse.ArgumentParser('DeiT evaluation script', add_help=False)
+
+    parser.add_argument('--model', default='TALL_SWIN', type=str)
+    parser.add_argument('--model_name', default="TALL_SWIN")
+    parser.add_argument('--batch-size', default=2, type=int)
+
+    # Dataset parameters
+    parser.add_argument('--data_txt_dir', type=str, default='##path_for_dataset_txt##')
+    parser.add_argument('--data_dir', type=str, default="##path_for_dataset##")
+    parser.add_argument('--dataset', default='ffpp', choices=list(DATASET_CONFIG.keys()))
+    parser.add_argument('--duration', default=1, type=int)
+    parser.add_argument('--frames_per_group', default=1, type=int)
+    parser.add_argument('--threed_data', default=False)
+    parser.add_argument('--input_size', default=224, type=int)
+    parser.add_argument('--disable_scaleup', action='store_true')
+    parser.add_argument('--random_sampling', action='store_true')
+    parser.add_argument('--dense_sampling', default=True)
+    parser.add_argument('--augmentor_ver', default='v1', type=str, choices=['v1', 'v2'])
+    parser.add_argument('--scale_range', default=[256, 320], type=int, nargs="+")
+    parser.add_argument('--modality', default='rgb', type=str)
+    parser.add_argument('--use_lmdb', default=False)
+    parser.add_argument('--use_pyav', default=False)
+
+    # temporal module / model params
+    parser.add_argument('--pretrained', action='store_true', default=False)
+    parser.add_argument('--temporal_module_name', default=None, type=str,
+                        choices=['ResNet3d', 'TAM', 'TTAM', 'TSM', 'TTSM', 'MSA'])
+    parser.add_argument('--temporal_attention_only', action='store_true', default=False)
+    parser.add_argument('--no_token_mask', action='store_true', default=False)
+    parser.add_argument('--temporal_heads_scale', default=1.0, type=float)
+    parser.add_argument('--temporal_mlp_scale', default=1.0, type=float)
+    parser.add_argument('--rel_pos', action='store_true', default=False)
+    parser.add_argument('--temporal_pooling', type=str, default=None,
+                        choices=['avg', 'max', 'conv', 'depthconv'])
+    parser.add_argument('--bottleneck', default=None, choices=['regular', 'dw'])
+
+    parser.add_argument('--window_size', default=7, type=int)
+    parser.add_argument('--thumbnail_rows', default=3, type=int)
+    parser.add_argument('--hpe_to_token', default=False, action='store_true')
+
+    parser.add_argument('--drop', type=float, default=0.0)
+    parser.add_argument('--drop-path', type=float, default=0.1)
+    parser.add_argument('--drop-block', type=float, default=None)
+
+    # runtime
+    parser.add_argument('--output_dir', default="./output")
+    parser.add_argument('--device', default='cuda')
+    parser.add_argument('--seed', default=42, type=int)
+    parser.add_argument('--num_workers', default=8, type=int)
+
+    parser.add_argument('--num_crops', default=1, type=int, choices=[1, 3, 5, 10])
+    parser.add_argument('--num_clips', default=3, type=int)
+
+    parser.add_argument('--world_size', default=1, type=int)
+    parser.add_argument("--local_rank", type=int)
+    parser.add_argument('--dist_url', default='env://')
+
+    # checkpoint
+    parser.add_argument('--initial_checkpoint', type=str, default='',
+                        help='path to .pth/.pth.tar checkpoint (expects key "model")')
+
+    parser.add_argument('--threshold', type=float, default=0.5,
+                        help='threshold to decide class 1 (fake) from prob[:,1]')
+    parser.add_argument('--metrics_out', default='', type=str,
+                        help='folder to save metrics.json and plots (default: output_dir)')
+    parser.add_argument('--save_plots', action='store_true',
+                        help='save cm.png / roc.png / pr.png')
+
+    return parser
+
+
+@torch.no_grad()
+def eval_with_outputs(data_loader, model, device, threshold: float = 0.5):
+    model.eval()
+    y_true, y_score, y_pred = [], [], []
+
+    thr = float(threshold)
+
+    for samples, targets in data_loader:
+        samples = samples.to(device, non_blocking=True)
+        targets = targets.to(device, non_blocking=True)
+
+        logits = model(samples)  # [B,2] or [B*K,2]
+
+        # if logits came per-clip, aggregate per video
+        B = targets.shape[0]
+        if logits.shape[0] != B:
+            if logits.shape[0] % B != 0:
+                raise RuntimeError(
+                    f"logits batch ({logits.shape[0]}) is not a multiple of target batch ({B})."
+                )
+            K = logits.shape[0] // B
+            logits = logits.view(B, K, -1).mean(dim=1)  # [B,2]
+
+        probs = torch.softmax(logits, dim=1)  # [B,2]
+        p1 = probs[:, 1]                      # class 1 (fake) score
+
+        # >>> THIS is the THRESHOLD <<<
+        hat = (p1 >= thr).long()
+
+        y_true.append(targets.detach().cpu().numpy())
+        y_score.append(p1.detach().cpu().numpy())
+        y_pred.append(hat.detach().cpu().numpy())
+
+    y_true = np.concatenate(y_true).astype(int)
+    y_score = np.concatenate(y_score).astype(float)
+    y_pred = np.concatenate(y_pred).astype(int)
+    return y_true, y_score, y_pred
+
+
+def plot_confusion(cm, out_path):
+    plt.figure(figsize=(6, 5))
+    plt.imshow(cm)
+    plt.title("Confusion Matrix")
+    plt.xlabel("Predicted")
+    plt.ylabel("True")
+    for (i, j), v in np.ndenumerate(cm):
+        plt.text(j, i, str(v), ha="center", va="center")
+    plt.tight_layout()
+    plt.savefig(out_path, dpi=200)
+    plt.close()
+
+
+def plot_roc(y, scores, out_path):
+    fpr, tpr, _ = roc_curve(y, scores)
+    auc = roc_auc_score(y, scores)
+    plt.figure(figsize=(7, 6))
+    plt.plot(fpr, tpr, label=f"AUC={auc:.4f}")
+    plt.plot([0, 1], [0, 1], "--", label="Chance")
+    plt.xlabel("FPR")
+    plt.ylabel("TPR")
+    plt.legend(loc="best")
+    plt.tight_layout()
+    plt.savefig(out_path, dpi=200)
+    plt.close()
+
+
+def plot_pr(y, scores, out_path):
+    p, r, _ = precision_recall_curve(y, scores)
+    ap = average_precision_score(y, scores)
+    plt.figure(figsize=(7, 6))
+    plt.plot(r, p, label=f"AP={ap:.4f}")
+    plt.xlabel("Recall")
+    plt.ylabel("Precision")
+    plt.legend(loc="best")
+    plt.tight_layout()
+    plt.savefig(out_path, dpi=200)
+    plt.close()
+
+
+def main(args):
+    utils.init_distributed_mode(args)
+    print(args)
+
+    device = torch.device(args.device)
+
+    seed = args.seed + utils.get_rank()
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    cudnn.benchmark = True
+
+    num_classes, train_list_name, val_list_name, test_list_name, filename_seperator, image_tmpl, filter_video, label_file = \
+        get_dataset_config(args.dataset, args.use_lmdb)
+
+    args.num_classes = num_classes
+    args.input_channels = 3 if args.modality == 'rgb' else 2 * 5
+
+    print(f"Creating model: {args.model}")
+    model = create_model(
+        args.model,
+        pretrained=args.pretrained,
+        duration=args.duration,
+        hpe_to_token=args.hpe_to_token,
+        rel_pos=args.rel_pos,
+        window_size=args.window_size,
+        thumbnail_rows=args.thumbnail_rows,
+        token_mask=not args.no_token_mask,
+        online_learning=False,
+        num_classes=args.num_classes,
+        drop_rate=args.drop,
+        drop_path_rate=args.drop_path,
+        drop_block_rate=args.drop_block,
+        use_checkpoint=False
+    )
+    model.to(device)
+
+    # mean/std
+    if args.distributed:
+        mean = (0.5, 0.5, 0.5) if 'mean' not in model.module.default_cfg else model.module.default_cfg['mean']
+        std = (0.5, 0.5, 0.5) if 'std' not in model.module.default_cfg else model.module.default_cfg['std']
+    else:
+        mean = (0.5, 0.5, 0.5) if 'mean' not in model.default_cfg else model.default_cfg['mean']
+        std = (0.5, 0.5, 0.5) if 'std' not in model.default_cfg else model.default_cfg['std']
+
+    # dataset (validation list)
+    video_data_cls = VideoDataSet
+    val_list = os.path.join(args.data_txt_dir, val_list_name)
+
+    val_augmentor = get_augmentor(
+        False, args.input_size, mean, std, args.disable_scaleup,
+        threed_data=args.threed_data, version=args.augmentor_ver,
+        scale_range=args.scale_range, num_clips=args.num_clips,
+        num_crops=args.num_crops, dataset=args.dataset
+    )
+
+    dataset_val = video_data_cls(
+        args.data_dir, val_list,
+        args.duration, args.frames_per_group,
+        num_clips=args.num_clips,
+        modality=args.modality,
+        dense_sampling=args.dense_sampling,
+        image_tmpl=image_tmpl,
+        transform=val_augmentor,
+        is_train=False, test_mode=False,
+        seperator=filename_seperator, filter_video=filter_video
+    )
+
+    data_loader_val = build_dataflow(
+        dataset_val, is_train=False, batch_size=args.batch_size,
+        workers=args.num_workers, is_distributed=args.distributed
+    )
+
+    if not args.initial_checkpoint:
+        raise RuntimeError("Please pass --initial_checkpoint pointing to the model checkpoint.")
+
+    checkpoint = torch.load(args.initial_checkpoint, map_location='cpu')
+    # many checkpoints come as {"model": state_dict, ...}
+    if isinstance(checkpoint, dict) and "model" in checkpoint:
+        utils.load_checkpoint(model, checkpoint["model"])
+    else:
+        # if it is a direct state_dict
+        model.load_state_dict(checkpoint, strict=False)
+
+    # eval
+    y_true, y_score, y_pred = eval_with_outputs(
+        data_loader_val, model, device, threshold=args.threshold
+    )
+
+    acc = accuracy_score(y_true, y_pred)
+    bacc = balanced_accuracy_score(y_true, y_pred)
+    prec, rec, f1, _ = precision_recall_fscore_support(
+        y_true, y_pred, average="binary", zero_division=0
+    )
+    cm = confusion_matrix(y_true, y_pred)
+
+    roc_auc = roc_auc_score(y_true, y_score)
+    pr_auc = average_precision_score(y_true, y_score)
+
+    print(f"\nN={len(y_true)} | thr={args.threshold:.3f}")
+    print(f"acc={acc:.4f} | bacc={bacc:.4f} | prec={prec:.4f} | rec={rec:.4f} | f1={f1:.4f} | roc_auc={roc_auc:.4f} | pr_auc={pr_auc:.4f}")
+    print(classification_report(y_true, y_pred, digits=4, zero_division=0))
+
+    outdir = args.metrics_out.strip() if args.metrics_out else args.output_dir
+    os.makedirs(outdir, exist_ok=True)
+
+    out_json = {
+        "threshold": float(args.threshold),
+        "acc": float(acc),
+        "balanced_acc": float(bacc),
+        "precision": float(prec),
+        "recall": float(rec),
+        "f1": float(f1),
+        "roc_auc": float(roc_auc),
+        "pr_auc": float(pr_auc),
+        "confusion_matrix": cm.tolist(),
+        "n": int(len(y_true)),
+    }
+    with open(os.path.join(outdir, "metrics.json"), "w", encoding="utf-8") as f:
+        json.dump(out_json, f, indent=2)
+
+    np.savez(os.path.join(outdir, "eval_outputs.npz"),
+             y_true=y_true, y_score=y_score, y_pred=y_pred)
+
+    if args.save_plots:
+        plot_confusion(cm, os.path.join(outdir, "cm.png"))
+        plot_roc(y_true, y_score, os.path.join(outdir, "roc.png"))
+        plot_pr(y_true, y_score, os.path.join(outdir, "pr.png"))
+        print(f"\n✔ Plots + metrics saved in: {os.path.abspath(outdir)}")
+    else:
+        print(f"\n✔ Metrics saved in: {os.path.abspath(os.path.join(outdir, 'metrics.json'))}")
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser('DeiT evaluation script', parents=[get_args_parser()])
+    args = parser.parse_args()
+    if args.output_dir:
+        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+    main(args)