scores_utils.py

#!/usr/bin/env python

from itertools import combinations
import json
from typing import Literal
from joblib import Parallel, delayed
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from sklearn.linear_model import LogisticRegression, SGDClassifier
from sklearn.model_selection import train_test_split
import torch
import os
from tqdm import tqdm
from sklearn.svm import SVC
from sklearn.decomposition import PCA
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis as LDA

from datasets import ASVspoof5
from trainers.utils import calculate_EER, calculate_minDCF

from config import local_config


def draw_score_distribution(c="FFConcat1", ep=15):
    # Load the scores
    scores_headers = ["AUDIO_FILE_NAME", "SCORE", "LABEL"]
    # scores_df = pd.read_csv(f"./scores/DF21/{c}_{c}_{ep}.pt_scores.txt", sep=",", names=scores_headers)
    scores_df = pd.read_csv(f"./scores/InTheWild/InTheWild_{c}_scores.txt", sep=",", names=scores_headers)

    # Filter the scores based on label
    bf_hist, bf_edges = np.histogram(scores_df[scores_df["LABEL"] == 0]["SCORE"], bins=15)
    sp_hist, sp_edges = np.histogram(scores_df[scores_df["LABEL"] == 1]["SCORE"], bins=15)
    bf_freq = bf_hist / np.sum(bf_hist)
    sp_freq = sp_hist / np.sum(sp_hist)
    bf_width = np.diff(bf_edges)
    sp_width = np.diff(sp_edges)
    plt.figure(figsize=(8, 5))
    plt.bar(
        bf_edges[:-1],
        bf_freq,
        width=(bf_width + sp_width) / 2,
        alpha=0.5,
        label="Bonafide",
        color="green",
        edgecolor="darkgreen",
        linewidth=1.5,
        align="edge",
    )
    plt.bar(
        sp_edges[:-1],
        sp_freq,
        width=(bf_width + sp_width) / 2,
        alpha=0.5,
        label="Spoofed",
        color="red",
        edgecolor="darkred",
        linewidth=1.5,
        align="edge",
    )
    plt.axvline(x=0.5, color="black", linestyle="--", label="Threshold 0.5", ymax=0.8, alpha=0.7)
    plt.xlabel("Scores - Probabilities of bonafide sample")
    plt.ylabel("Relative frequency of bonafide/spoofed samples")
    plt.title(f"Distribution of scores: {c}")
    plt.legend(loc="upper center")
    # plt.xlim(0, 1)
    plt.savefig(f"./scores/{c}_{ep}_scores.png")


def draw_det(dataset: Literal["DF21", "InTheWild"], c="FFLSTM", ep=10):
    # Load the scores
    scores_headers = ["AUDIO_FILE_NAME", f"SCORE_{c}", "LABEL"]
    name = f"{c}_{c}_{ep}.pt_scores.txt" if dataset == "DF21" else f"InTheWild_{c}_scores.txt"
    scores_df = pd.read_csv(f"./scores/{dataset}/{name}", sep=",", names=scores_headers)
    calculate_EER(c, scores_df["LABEL"], scores_df[f"SCORE_{c}"], True, f"{dataset}_{c}")


def split_scores_VC_TTS(filename: str):
    # Load the scores
    scores_headers = ["AUDIO_FILE_NAME", "SCORE", "LABEL"]
    scores_df = pd.read_csv(filename, sep=",", names=scores_headers)
    scores_df["SCORE"] = scores_df["SCORE"].astype(float)
    print(f"Processing {filename}")

    # Load DF21 protocol
    la19_headers = ["SPEAKER_ID", "AUDIO_FILE_NAME", "-", "MODIF", "KEY"]
    protocol_df = pd.read_csv(f"ASVspoof2019.LA.cm.eval.trl.txt", sep=" ")
    protocol_df.columns = la19_headers
    protocol_df = protocol_df.merge(scores_df, on="AUDIO_FILE_NAME")
    eer = calculate_EER(filename, protocol_df["LABEL"], protocol_df["SCORE"], False, f"LA19")
    print(f"EER for LA19: {eer*100:.3f}%")

    asvspoof_bonafide_df = protocol_df[(protocol_df["KEY"] == "bonafide")].reset_index(drop=True)

    tts_systems = ["A01", "A02", "A03", "A04", "A07", "A08", "A09", "A10", "A11", "A12", "A16"]
    tts_subset = protocol_df[protocol_df["MODIF"].isin(tts_systems)].reset_index(drop=True)
    tts_subset = pd.concat([tts_subset, asvspoof_bonafide_df], axis=0)

    vc_systems = ["A05", "A06", "A17", "A18", "A19"]
    vc_subset = protocol_df[protocol_df["MODIF"].isin(vc_systems)].reset_index(drop=True)
    # vcc_subset = protocol_df[protocol_df["SOURCE"].str.contains("vcc")].reset_index(drop=True)
    vc_subset = pd.concat([vc_subset, asvspoof_bonafide_df], axis=0)

    for subset, subset_df in zip(["TTS ", "VC  "], [tts_subset, vc_subset]):
        eer = calculate_EER(filename, subset_df["LABEL"], subset_df["SCORE"], False, f"{subset}")
        print(f"EER for {subset}: {eer*100:.3f}%")


def split_scores_asvspoof_VCC(filename):
    # Load the scores
    scores_headers = ["AUDIO_FILE_NAME", "SCORE", "LABEL"]
    scores_df = pd.read_csv(filename, sep=",", names=scores_headers)
    scores_df["SCORE"] = scores_df["SCORE"].astype(float)
    print(f"Processing {filename}")

    # Load DF21 protocol
    df21_headers = [
        "SPEAKER_ID",
        "AUDIO_FILE_NAME",
        "-",
        "SOURCE",
        "MODIF",
        "KEY",
        "-",
        "VARIANT",
        "-",
        "-",
        "-",
        "-",
        "-",
    ]
    protocol_df = pd.read_csv("trial_metadata.txt", sep=" ")
    protocol_df.columns = df21_headers
    protocol_df = protocol_df.merge(scores_df, on="AUDIO_FILE_NAME")
    eer = calculate_EER(filename, protocol_df["LABEL"], protocol_df["SCORE"], False, f"DF21")
    print(f"EER for DF21: {eer*100:.3f}%")

    asvspoof_subset = protocol_df[protocol_df["SOURCE"].str.contains("asvspoof")].reset_index(drop=True)
    vcc_subset = protocol_df[protocol_df["SOURCE"].str.contains("vcc")].reset_index(drop=True)

    for i, (subset, subset_df) in enumerate(zip(["asvs", "vcc "], [asvspoof_subset, vcc_subset])):
        eer = calculate_EER(filename, subset_df["LABEL"], subset_df["SCORE"], False, f"{subset}")
        print(f"EER for {subset}: {eer*100:.3f}%")


def get_all_scores_df(variant: Literal["DF21", "InTheWild"]) -> pd.DataFrame:
    all_scores_df = pd.DataFrame()
    for c, ep in [
        ("FFDiff", 20),
        ("FFDiffAbs", 15),
        ("FFDiffQuadratic", 15),
        ("FFConcat1", 15),
        ("FFConcat2", 10),
        ("FFConcat3", 10),
        # ("FFLSTM", 10),
        ("FFLSTM2", 15),
    ]:
        print(f"Loading scores for {c}_{ep}")
        scores_headers = ["AUDIO_FILE_NAME", f"SCORE_{c}", "LABEL"]
        scores_df = pd.read_csv(
            f"./scores/{variant}/{(c+'_'+c+'_'+str(ep)+'.pt_scores.txt') if variant == 'DF21' else 'InTheWild_'+c+'_scores.txt'}",
            sep=",",
            names=scores_headers,
        )
        if all_scores_df.empty:
            all_scores_df.insert(0, "AUDIO_FILE_NAME", scores_df["AUDIO_FILE_NAME"])
            all_scores_df.insert(1, "LABEL", scores_df["LABEL"])
            all_scores_df.insert(2, f"SCORE_{c}", scores_df[f"SCORE_{c}"])
        else:
            all_scores_df = all_scores_df.merge(scores_df, on=["AUDIO_FILE_NAME", "LABEL"])
    return all_scores_df


def fusion_NN(variant: Literal["DF21", "InTheWild"]):
    d = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    layer = torch.nn.Linear(7, 2).to(d)

    all_scores_df = get_all_scores_df(variant)

    batch_size = 1280  # local 1060 has 1280 CUDA cores
    loss_fn = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(layer.parameters())
    for epoch in range(1, 201):  # 1-indexed epochs
        names = []
        losses = []
        probs = []
        accuracies = []
        labels = []
        for i in tqdm(range(0, len(all_scores_df), batch_size)):
            optimizer.zero_grad()

            batch_scores = torch.tensor(
                all_scores_df.iloc[i : i + batch_size, 2:].values, dtype=torch.float32
            ).to(d)
            batch_labels = torch.tensor(
                all_scores_df.iloc[i : i + batch_size, 1].values, dtype=torch.long
            ).to(d)

            outputs = layer(batch_scores)
            probs_batch = torch.nn.functional.softmax(outputs, dim=1)

            loss = loss_fn(outputs, batch_labels)

            names.extend(all_scores_df.iloc[i : i + batch_size, 0].tolist())
            losses.extend(loss.item() for _ in range(batch_size))
            probs.extend(probs_batch[:, 0].tolist())
            accuracies.extend((torch.argmax(probs_batch, 1) == batch_labels).tolist())
            labels.extend(batch_labels.tolist())

            loss.backward()
            optimizer.step()

        eer = calculate_EER("Fusion", labels, probs, False, "Fusion")
        print(f"Epoch {epoch}: Loss: {np.mean(losses)}, Acc: {np.mean(accuracies)*100}%, EER: {eer*100}%")
        print(f"Estimated parameters: {layer.weight}, {layer.bias}")

        with open(f"./scores/{variant}/fusion/NN_{epoch}_scores.txt", "w") as f:
            for file_name, score, label in zip(names, probs, labels):
                f.write(f"{file_name},{score},{int(label)}\n")


def fusion_scores(dataset: Literal["DF21", "ITW", "ASVspoof5", "AS5dev"], adversarial_only: bool = False, adversarial_inverted: bool = False):
    # region old
    # dfs = []

    # for file in os.listdir(f"./scores/{dataset}"):
    #     if ".json" in file or "fusion" in file:
    #         continue  # Skip the fusion scores

    #     df = pd.read_csv(
    #         f"./scores/{dataset}/{file}", header=None, names=["file", f'score_{file.split("_")[1]}', "label"]
    #     )
    #     dfs.append(df)

    # final_df = pd.concat(dfs, axis=0).groupby(["file", "label"]).first().reset_index()
    # if dataset == "DF21":
    #     final_df = final_df.drop(columns=["score_FF"])  # Only pair-input systems

    # scores = [
    #     "score_FFConcat1",
    #     "score_FFConcat2",
    #     "score_FFConcat3",
    #     "score_FFDiff",
    #     "score_FFDiffAbs",
    #     "score_FFDiffQuadratic",
    #     # "score_FFLSTM",
    #     "score_FFLSTM2",
    # ]
    # comb = []
    # for i in range(2, len(scores) + 1):
    #     comb.extend(combinations(scores, i))

    # score_dict = {}
    # for combination in tqdm(comb):
    #     name = f" + ".join(combination)
    #     mean_score = final_df[list(combination)].mean(axis=1)
    #     max_score = final_df[list(combination)].max(axis=1)
    #     min_score = final_df[list(combination)].min(axis=1)
    #     sqrt_score = final_df[list(combination)].apply(lambda x: x.prod() ** (1 / len(combination)), axis=1)

    #     mean_eer = calculate_EER(name, final_df["label"], mean_score, False, "")
    #     max_eer = calculate_EER(name, final_df["label"], max_score, False, "")
    #     min_eer = calculate_EER(name, final_df["label"], min_score, False, "")
    #     sqrt_eer = calculate_EER(name, final_df["label"], sqrt_score, False, "")

    #     score_dict[name] = {"mean": mean_eer, "max": max_eer, "min": min_eer, "sqrt": sqrt_eer}

    # json.dump(score_dict, open(f"./scores/{dataset}/fusion_scores.json", "w"))
    # endregion

    dfs = []
    for file in os.listdir(f"./scores/final/"):
        if dataset not in file or ".json" in file:
            continue

        df = pd.read_csv(
            f"./scores/final/{file}",
            header=None,
            names=["file", f'score_{file.split("_")[0]}_{file.split("_")[1]}', "label"],
        )

        dfs.append(df)

    final_df = pd.concat(dfs, axis=0).groupby(["file", "label"]).first().reset_index()
    # final_df = final_df.drop(columns=["score_FF_MHFA", "score_FF_AASIST"])  # Only pair-input systems
    # final_df = final_df.drop(columns=[col for col in final_df.columns if "Diff" in col])  # Drop diff models as they suck

    if adversarial_only:
        # Filter to keep only adversarial recordings
        adversarial_recordings = asvspoof5_extract_adversarial_recording_names(inverted=adversarial_inverted)
        final_df = final_df[final_df["file"].isin(adversarial_recordings)].reset_index(drop=True)

    # do combinations from all the score_* columns
    scores = [col for col in final_df.columns if col.startswith("score_FF")]
    comb = []
    # for i in range(2, 3):
    # for i in range(2, len(scores) + 1):
    # for i in range(8, 9):
    # comb.extend(combinations(scores, 1)) # single
    # comb.extend(combinations(scores, 2)) # pairs
    # comb.extend(combinations(scores, len(scores))) # all

    # for i in range(2, len(scores) + 1):
    # for i in range(1, 2):
    #     comb.extend(combinations(scores, i))
    comb.extend([("score_FFAttn2_MHFA", "score_FFAttn2_AASIST")])
    # comb.extend(combinations(scores, 17))
    # comb.extend(combinations(scores, 16))
    # comb.extend(combinations(scores, 15))
    # comb.extend(combinations(scores, 14))

    score_dict = {}
    # for combination in tqdm(comb):
    def evaluate_combination(combination):
        name = f" + ".join(combination)
        # mean_score = final_df[list(combination)].mean(axis=1)
        # max_score = final_df[list(combination)].max(axis=1)
        # min_score = final_df[list(combination)].min(axis=1)
        # sqrt_score = (final_df[list(combination)] ** 0.5).sum(axis=1) / len(combination)
        # geom_score = final_df[list(combination)].prod(axis=1) ** (1 / len(combination))

        lr = LogisticRegression()
        lr.fit(final_df[list(combination)], final_df["label"])
        scores = lr.predict_proba(final_df[list(combination)])

        # mean_eer = calculate_EER(name, final_df["label"], mean_score, False, "")
        # max_eer = calculate_EER(name, final_df["label"], max_score, False, "")
        # min_eer = calculate_EER(name, final_df["label"], min_score, False, "")
        # sqrt_eer = calculate_EER(name, final_df["label"], sqrt_score, False, "")
        # geom_eer = calculate_EER(name, final_df["label"], geom_score, False, "")

        # lr_eer = calculate_EER(name, final_df["label"], scores[:, 0], False, "")
        # print("Combination:", name, "LR EER:", lr_eer*100, "%")
        lr_mindcf = calculate_minDCF(final_df["label"], scores[:, 0], p_target=0.95, c_miss=1, c_fa=10)
        print(f"Combination: {name}, LR minDCF: {lr_mindcf}")

        score_dict[name] = {
            # "mean": mean_eer,
            # "max": max_eer,
            # "min": min_eer,
            # "sqrt": sqrt_eer,
            # "geom": geom_eer,
            "lr": lr_mindcf,
        }

    Parallel(n_jobs=-1, backend="threading")(delayed(evaluate_combination)(combination) for combination in tqdm(comb))

    json.dump(score_dict, open(f"./scores/final/fusion_scores_{dataset}_mindcf{'_adversarial' if (adversarial_only and not adversarial_inverted) else ('_notadversarial' if (adversarial_only and adversarial_inverted) else '')}.json", "w"))


def fusion_scores_from_json(
    dataset: Literal["DF21", "InTheWild"], number: Literal["all", "oneplusone"] = "all"
):
    scores = json.load(open(f"./scores/{dataset}/fusion_scores.json", "r"))

    if number == "oneplusone":
        doubles = {key: scores[key] for key in scores if len(key.split(" + ")) == 2}
        scores = {
            key: doubles[key]
            for key in doubles
            if (key.count("FFDiff") == 1 and (key.count("FFConcat") == 1 or key.count("FFLSTM") == 1))
        }

    for fusion in ["mean", "max", "min", "sqrt"]:
        best_fusion = min(scores, key=lambda x: scores[x][fusion])
        print(f"Best {fusion} fusion: {best_fusion}, EER: {scores[best_fusion][fusion]*100}%")


def fusion_LDA(variant: Literal["DF21", "InTheWild"]):
    all_scores_df = get_all_scores_df(variant)

    lda = LDA()
    lda.fit(all_scores_df.iloc[:, 2:], all_scores_df["LABEL"])
    x_trans = lda.transform(all_scores_df.iloc[:, 2:])

    X_train, X_test, y_train, y_test = train_test_split(x_trans, all_scores_df["LABEL"], test_size=0.8)

    for k in ["linear", "poly", "rbf", "sigmoid"]:
        svm = SVC(kernel=k, probability=True)
        svm.fit(X_train, y_train)
        scores = svm.predict_proba(X_test)

        eer = calculate_EER(f"LDA_{k}", y_test, scores[:, 0], False, f"LDA_{k}")
        print(f"LDA + SVM({k}) EER: {eer*100}%")


def fusion_PCA(variant: Literal["DF21", "InTheWild"]):
    all_scores_df = get_all_scores_df(variant)

    pca = PCA()
    pca.fit(all_scores_df.iloc[:, 2:])
    x_trans = pca.transform(all_scores_df.iloc[:, 2:])

    X_train, X_test, y_train, y_test = train_test_split(x_trans, all_scores_df["LABEL"], test_size=0.8)

    for k in ["linear", "poly", "rbf", "sigmoid"]:
        svm = SVC(kernel=k, probability=True)
        svm.fit(X_train, y_train)
        scores = svm.predict_proba(X_test)

        eer = calculate_EER(f"PCA_{k}", y_test, scores[:, 0], False, f"PCA_{k}")
        print(f"PCA + SVM({k}) EER: {eer*100}%")


def asvspoof5_extract_adversarial_recording_names(inverted: bool = False) -> pd.Series:
    protocol_headers = [
        "SPEAKER_ID",
        "AUDIO_FILE_NAME",
        "GENDER",
        "CODEC",
        "CODEC_Q",
        "CODEC_SEED",
        "ATTACK_TAG",
        "ATTACK_LABEL",
        "KEY",
        "-",
    ]

    protocol_df = pd.read_csv("./ASVspoof5.eval.track_1.tsv", sep=" ", header=None)
    protocol_df.columns = protocol_headers

    # Select only the adversarial attacks
    if inverted:
        protocol_df = protocol_df[~protocol_df["ATTACK_LABEL"].isin(["A18", "A20", "A23", "A27", "A30", "A31", "A32"])].reset_index(drop=True)
    else:
        protocol_df = protocol_df[protocol_df["ATTACK_LABEL"].isin(["bonafide", "A18", "A20", "A23", "A27", "A30", "A31", "A32"])].reset_index(drop=True)

    return protocol_df["AUDIO_FILE_NAME"]


def calculate_stdev():
    scores = {
        "FFAttn1_MHFA": [4.457, 4.217, 4.250],
        "FFAttn2_AASIST": [3.681, 3.699, 3.685],
        "FFAttn2_MHFA": [3.326, 3.370, 3.377],
    }

    stdev_sum = 0
    for key, eers in scores.items():
        mean_eer = np.mean(eers)
        stdev_sum += sum((eer - mean_eer) ** 2 for eer in eers)
        print(f"{key} mean EER: {mean_eer}%, stdev: {np.std(eers):.4f}%")
    global_stdev = (stdev_sum / sum(len(eers) for eers in scores.values())) ** 0.5
    print(f"Global stdev: {global_stdev:.4f}%")

if __name__ == "__main__":
    # calculate eer for all score files in the scores directory
    # for file in os.listdir("./scores"):
    #     if not file.endswith("scores.txt"):
    #         continue

    #     print(f"Calculating EER for {file}:", end=" ")
    #     scores_headers = ["AUDIO_FILE_NAME", "SCORE", "LABEL"]
    #     scores_df = pd.read_csv(f"./scores/{file}", sep=",", names=scores_headers)
    #     eer = calculate_EER(file, scores_df["LABEL"], scores_df["SCORE"], False, file)
    #     print(f"{eer*100:.3}%")

    for dataset in ["ASVspoof5"]:
        print("================ Fusion scores for", dataset, "================")
        # LR fusion

        # Standard fusion
        # fusion_scores(dataset, adversarial_only=False, adversarial_inverted=False)

        """
        scores = json.load(open(f"./scores/final/fusion_scores_{dataset}_single.json", "r"))
        # keep only the combinations of scores that have both AASIST and MHFA
        scores = {
            key: scores[key]
            for key in scores
            # if ("AASIST" in key and "MHFA" in key) and key.count(" + ") == 1
        }

        best_fusion = {}
        print("================================================================")
        # for fusion in ["mean", "max", "min", "sqrt", "geom", "lr"]:
        for fusion in ["lr"]:
            min_fusion = min(scores, key=lambda x: scores[x][fusion])
            print(f"Best {dataset} {fusion} fusion: {min_fusion}, EER: {scores[min_fusion][fusion]*100}%")
            best_fusion[fusion] = (min_fusion, scores[min_fusion][fusion] * 100)

        bf = min(best_fusion, key=lambda x: best_fusion[x][1])
        print(f">>>>> Best {dataset} fusion: {bf} ({best_fusion[bf][0]}), EER: {best_fusion[bf][1]}% <<<<<")
        """

    # LA19
    # for file in os.listdir("./scores/final/"):
    #     if "LA19" in file:
    #         split_scores_VC_TTS(f"./scores/final/{file}")
    #         print("#####" * 10)

    # DF21
    # for file in os.listdir("./scores/final/"):
    #     if "DF21" in file and ".json" not in file:
    #         split_scores_asvspoof_VCC(f"./scores/final/{file}")
    #         print("#####" * 10)