eval_closedset_id.py

import os

import argparse
import torch
from torchvision import models
import torchvision.transforms.functional as F
from torchvision import transforms
import torch.nn as nn
import numpy as np
import pandas as pd
from scipy.stats import entropy
from sklearn.metrics import confusion_matrix, balanced_accuracy_score
import matplotlib.pyplot as plt
import seaborn as sns

from matplotlib.colors import LinearSegmentedColormap

from datasets import PrecomputedCropDataset



def build_model(num_classes, backbone="resnet18"):
    if backbone == "resnet18":
        model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
        #model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
        in_features = model.fc.in_features
        model.fc = nn.Linear(in_features, num_classes)

    elif backbone == "resnet50":
        model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
        in_features = model.fc.in_features
        model.fc = nn.Linear(in_features, num_classes)

    elif backbone == "resnet101":
        model = models.resnet101(weights=models.ResNet101_Weights.IMAGENET1K_V1)
        in_features = model.fc.in_features
        model.fc = nn.Linear(in_features, num_classes)

    else:
        raise ValueError(f"Unknown backbone: {backbone}")

    return model

def aggregate_track(probs, logits, method="mean_probs", threshold1=0.5, threshold2=0.9, k=0.5):
    if method == "mean_logits":
        pred = logits.mean(axis=0).argmax()

    elif method == "mean_probs":
        pred = probs.mean(axis=0).argmax()

    elif method == "max_probs":
        pred = probs.max(axis=0).argmax()

    elif method == "threshold":
        counts = (probs > threshold1).sum(axis=0)

        if counts.sum() == 0:
            pred = probs.mean(axis=0).argmax()
        else:
            max_count = counts.max()
            candidates = np.where(counts == max_count)[0]
            if len(candidates) == 1:
                pred = candidates[0]
            else:
                totals = probs.sum(axis=0)
                pred = candidates[np.argmax(totals[candidates])]

    elif method == "double_threshold":
        weights = (
            (probs > threshold1).astype(float) * 1.0 +
            (probs > threshold2).astype(float) * 100.0
        )

        scores = weights.sum(axis=0)

        if scores.sum() == 0:
            pred = probs.mean(axis=0).argmax()
        else:
            pred = scores.argmax()

    elif method == "exponential_weight":
        weights = np.exp(9.2 * (probs - 0.5))  # exponential weighting
        scores = weights.sum(axis=0)
        pred = scores.argmax()

    elif method == "entropy_weight":
        eps = 1e-6
        entropy = -np.sum(probs * np.log(probs + eps), axis=1)  # (T,)
        weights = 1.0 / (entropy + eps)                          # (T,)

        scores = (weights[:, None] * probs).sum(axis=0)
        pred = scores.argmax()

    elif method == "entropy_threshold":
        eps = 1e-6
        entropy = -np.sum(probs * np.log(probs + eps), axis=1)

        entropy_thr = np.percentile(entropy, 10*threshold1)  # or fixed value
        mask = entropy < entropy_thr

        if mask.sum() == 0:
            pred = probs.mean(axis=0).argmax()
        else:
            scores = probs[mask].sum(axis=0)
            pred = scores.argmax()

    elif method == "top_k_aggregation":  # keep top 50%

        conf = probs.max(axis=1)  # per-frame confidence
        num_keep = max(1, int(len(conf) * k))
        top_idx = np.argsort(conf)[-num_keep:]

        scores = probs[top_idx].sum(axis=0)
        pred = scores.argmax()
        
    
    return pred



def validate_and_save(model, loader, aggregate_method, device, output_prefix, k=0.5, threshold1=0.5, 
                      threshold2=0.9, write_frame_preds=False, write_track_preds=False):
    model.eval()

    records = []

    # --------------------------------------------------
    # Frame-level inference + saving
    # --------------------------------------------------
    with torch.no_grad():
        for imgs, experiment, track_id, labels, frame_nums in loader:
            imgs = imgs.to(device)
            labels = labels.to(device)

            logits = model(imgs)
            probs = logits.softmax(dim=1)


            logits_np = logits.cpu().numpy()
            probs_np = probs.cpu().numpy()
            labels_np = labels.cpu().numpy()
            track_ids_np = track_id.cpu().numpy()
            frame_nums_np = frame_nums.cpu().numpy()

            for i in range(len(imgs)):
                records.append({
                    "experiment": experiment[i],
                    "track_id": int(track_ids_np[i]),   # CRITICAL
                    "frame_num": int(frame_nums_np[i]),
                    "label": int(labels_np[i]),
                    "logits": logits_np[i],
                    "probs": probs_np[i],
                })

    df = pd.DataFrame(records)

    if write_frame_preds:
        frame_csv = f"output/{output_prefix.split('_')[0]}/{output_prefix}_frame_results.csv"
        df.to_csv(frame_csv, index=False)
        #print(f"Saved frame-level results to {frame_csv}")

    # --------------------------------------------------
    # Track-level aggregation (MATCHES validate())
    # --------------------------------------------------
    track_results = []

    for (experiment, track_id), group in df.groupby(["experiment", "track_id"]):
        probs = np.stack(group["probs"].values)  # (T, C)
        logits = np.stack(group["logits"].values)
        labels = group["label"].values

        gt = int(labels[0])  # same for entire track

        pred = aggregate_track(probs, logits, method = aggregate_method, k=k, threshold1=threshold1, threshold2=threshold2)

        track_results.append({
            "experiment": experiment,
            "track_id": track_id,
            "label": gt,
            "pred": int(pred),
            "num_frames": len(group),
        })

    track_df = pd.DataFrame(track_results)


    if write_track_preds:
        track_csv = f"output/{output_prefix.split('_')[0]}/{output_prefix}_track_results.csv"
        track_df.to_csv(track_csv, index=False)
        print(f"Saved track-level results to {track_csv}")

    accuracy = (track_df["label"] == track_df["pred"]).mean()
    print(f"Track-wise accuracy: {accuracy:.4f}")

    # compute class-balanced accuracy (mean per-class recall)
    unique_labels = np.unique(track_df["label"].values)
    cm = confusion_matrix(track_df["label"].values, track_df["pred"].values, labels=unique_labels)
    with np.errstate(divide='ignore', invalid='ignore'):
        per_class_recall = np.diag(cm) / cm.sum(axis=1)
    per_class_recall = np.nan_to_num(per_class_recall, nan=0.0)
    valid = cm.sum(axis=1) > 0
    class_balanced_accuracy = float(per_class_recall[valid].mean())
    

    balanced_acc = balanced_accuracy_score(
        track_df["label"].values,
        track_df["pred"].values
    )

    # check that both methods give the same result
    assert np.isclose(class_balanced_accuracy, balanced_acc), \
        f"Mismatch in balanced accuracy: {class_balanced_accuracy} vs {balanced_acc}"
    
    #class_balanced_accuracy = float(per_class_recall.mean())
    print(f"Class-balanced accuracy: {class_balanced_accuracy:.4f}")

    return accuracy, class_balanced_accuracy, track_df





def main():
    
    parser = argparse.ArgumentParser()
    parser.add_argument("--val_root", required=True)
    parser.add_argument('--model_path', type=str, default="checkpoints/reid_resnet18_baseline_count_001/reid_resnet18_baseline_count_001_best.pt")
    parser.add_argument('--group', type=str, default="Alpha")
    parser.add_argument("--backbone", default="resnet18")
    parser.add_argument("--aggregate_method", type=str, default="mean_probs",
                        choices=["mean_logits", "mean_probs", "max_probs",
                                 "threshold", "double_threshold",
                                 "exponential_weight", "entropy_weight",
                                 "entropy_threshold", "top_k_aggregation"])
    parser.add_argument("--k", type=float, default=0.5)
    parser.add_argument("--threshold1", type=float, default=0.5)
    parser.add_argument("--threshold2", type=float, default=0.9)
    parser.add_argument("--write_frame_preds", action="store_true")
    parser.add_argument("--write_track_preds", action="store_true")
    parser.add_argument("--write_summary", action="store_true")
    parser.add_argument("--plot_conf_matrix", action="store_true")
    args = parser.parse_args()

    val_tf = transforms.Compose([
        #transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])

    val_ds = PrecomputedCropDataset(
        root_dir=args.val_root,
        id_labels=True,
        transform=val_tf
    )

    dataloader = torch.utils.data.DataLoader(val_ds, batch_size=64, shuffle=False, num_workers=4)

    num_classes = max(val_ds.labels) + 1

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    model = build_model(num_classes, args.backbone).to(device)
    #load weights
    model.load_state_dict(torch.load(args.model_path))
    
    os.makedirs(f"output/{args.group.split('_')[0]}", exist_ok=True)

    acc, cl_acc, track_df = validate_and_save(model, dataloader, args.aggregate_method, device, output_prefix=args.group, 
                                              k=args.k, threshold1=args.threshold1, threshold2=args.threshold2,
                                              write_frame_preds=args.write_frame_preds, write_track_preds=args.write_track_preds)


    if args.write_summary:
        # Write args.aggregate_method, args.group, acc, cl_acc to a summary file
        with open(f"output/{args.group.split('_')[0]}/summary.csv", "a") as f:
            f.write(f"{args.aggregate_method},{args.group},{acc},{cl_acc}\n")

    #make a connfusion matrix
    all_names = ['Floreana', 'Rabita', 'Genovesa', 'Redonda',
             'George', 'Hermanos', 'Pinta', 'Marchena']

        
    kept_labels = list(range(len(all_names)))
    kept_names = all_names

    if args.plot_conf_matrix:
        cm = confusion_matrix(track_df['label'], track_df['pred'], labels=kept_labels)

        # --- build a custom colormap:
        # off-diagonal: orange-ish, diagonal: blue-ish (by boosting diagonal values)
        #maxv = cm.max() if cm.size else 1
        #diag_boost = (maxv + 1)  # ensures diagonal stands out as "more intense"
        #cm_for_color = cm.astype(float).copy()
        #np.fill_diagonal(cm_for_color, np.diag(cm_for_color) + diag_boost)

        cmap = LinearSegmentedColormap.from_list(
            "orange_to_blue",
            ["#fff5eb", "#fdae6b", "#f16913"], #, "#9ecae1", "#3182bd", "#08519c"
            N=256
        )

        sns.set_theme(style="white", font_scale=1.25)  # slightly larger text

        fig, ax = plt.subplots(figsize=(11, 9), constrained_layout=True)
        sns.heatmap(
            cm,
            annot=cm, fmt="d",                 # show original counts as text
            cmap=cmap,
            cbar=False,
            square=True,
            linewidths=0.5, linecolor="white",
            xticklabels=kept_names, yticklabels=kept_names,
            annot_kws={"size": 12},
            ax=ax
        )

        ax.set_xlabel("Predicted", labelpad=10)
        ax.set_ylabel("True", labelpad=10)
        ax.set_title("Confusion matrix, ResNet50", pad=12, fontsize=16, weight="bold")

        # y labels horizontal; x labels slightly rotated
        ax.set_yticklabels(ax.get_yticklabels(), rotation=0, va="center")
        ax.set_xticklabels(ax.get_xticklabels(), rotation=25, ha="right")

        # optional: cleaner spines
        for spine in ax.spines.values():
            spine.set_visible(False)

        fig.savefig(f"output/{args.group.split('_')[0]}/{args.group}_confusion_matrix.png", dpi=300)
        plt.close(fig)



if __name__ == '__main__':
    torch.cuda.set_device(0)

    main()