scripts/calibrate_threshold_wav2vec.py

import os, argparse, numpy as np
from glob import glob

try: from app.inference_wav2vec import Detector
except ImportError: from inference_wav2vec import Detector

def collect(root):
    H = sorted(glob(os.path.join(root, "human", "*.wav")))
    A = sorted(glob(os.path.join(root, "ai", "*.wav")))
    return H, A

def main(a):
    det = Detector(weights_path=a.weights)
    H, A = collect(a.root)
    ys, ps_mic, ps_up = [], [], []
    for p in H:
        ys.append(0)
        ps_mic.append(det.predict_proba(p, source_hint="microphone")["ai"])
        ps_up.append(det.predict_proba(p, source_hint="upload")["ai"])
    for p in A:
        ys.append(1)
        ps_mic.append(det.predict_proba(p, source_hint="microphone")["ai"])
        ps_up.append(det.predict_proba(p, source_hint="upload")["ai"])
    ys = np.array(ys); ps_mic = np.array(ps_mic); ps_up = np.array(ps_up)

    def sweep(ps):
        best = (0.5, -1, 1e9)
        for thr in np.linspace(0.5, 0.8, 61):
            pred = (ps >= thr).astype(int)
            tp = ((pred==1)&(ys==1)).sum(); fp = ((pred==1)&(ys==0)).sum()
            fn = ((pred==0)&(ys==1)).sum()
            prec = tp / max(tp+fp,1); rec = tp / max(tp+fn,1)
            f1 = 2*prec*rec / max(prec+rec,1e-9)
            if (f1 > best[1]) or (f1==best[1] and fp < best[2]):
                best = (float(thr), float(f1), int(fp))
        return best

    mt, mf1, mfp = sweep(ps_mic)
    ut, uf1, ufp = sweep(ps_up)
    print(f"MIC threshold ~ {mt:.2f} (F1={mf1:.3f}, human_as_ai={mfp})")
    print(f"UPLOAD threshold ~ {ut:.2f} (F1={uf1:.3f}, human_as_ai={ufp})")
    print("Set:")
    print(f'  $env:DETECTOR_MIC_THRESHOLD="{mt:.2f}"')
    print(f'  $env:DETECTOR_UPLOAD_THRESHOLD="{ut:.2f}"')

if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", required=True, help="folder with human/ and ai/")
    ap.add_argument("--weights", default="app/models/weights/wav2vec2_classifier.pth")
    main(ap.parse_args())