Add evaluation scripts

Author: bshall

abx.py

@@ -0,0 +1,79 @@
+import argparse
+import ast
+from pathlib import Path
+
+import ABXpy.analyze as analyze
+import ABXpy.distances.distances as distances
+import ABXpy.distances.metrics.cosine as cosine
+import ABXpy.distances.metrics.dtw as dtw
+import ABXpy.score as score
+import pandas
+from ABXpy.misc.any2h5features import convert
+
+
+def dtw_cosine_distance(x, y, normalized):
+    return dtw.dtw(x, y, cosine.cosine_distance, normalized)
+
+
+def average_abx(filename, task_type):
+    df = pandas.read_csv(filename, sep='\t')
+    if task_type == "across":
+        # aggregate on context
+        groups = df.groupby(["speaker_1", "speaker_2", "phone_1", "phone_2"],
+                            as_index=False)
+        df = groups["score"].mean()
+    elif task_type == "within":
+        arr = list(map(ast.literal_eval, df["by"]))
+        df["speaker"] = [e for e, f, g in arr]
+        df["context"] = [f for e, f, g in arr]
+
+        # aggregate on context
+        groups = df.groupby(["speaker", "phone_1", "phone_2"], as_index=False)
+        df = groups["score"].mean()
+    else:
+        raise ValueError("Unknown task type: {0}".format(task_type))
+
+    # aggregate on talker
+    groups = df.groupby(["phone_1", "phone_2"], as_index=False)
+    df = groups['score'].mean()
+    average = df.mean()[0]
+    average = (1.0 - average) * 100
+    return average
+
+
+def evaluate_abx(args):
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    feature_path = out_dir / "features.features"
+    distance_path = out_dir / "data.distance"
+    score_path = out_dir / "data.score"
+    analyze_path = out_dir / "data.csv"
+
+    if not feature_path.exists():
+        convert(args.feature_dir, h5_filename=str(feature_path))
+
+    if not distance_path.exists():
+        distances.compute_distances(
+            str(feature_path), "features", str(args.task_path),
+            str(distance_path), dtw_cosine_distance,
+            normalized=True, n_cpu=6)
+
+    if not score_path.exists():
+        score.score(str(args.task_path), str(distance_path), str(score_path))
+
+    if not analyze_path.exists():
+        analyze.analyze(str(args.task_path), str(score_path), str(analyze_path))
+
+    abx = average_abx(str(analyze_path), args.task_type)
+    print("average abx: {:.3f}".format(abx))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--task-type", type=str)
+    parser.add_argument("--task-path", type=str)
+    parser.add_argument("--feature-dir", type=str)
+    parser.add_argument("--out-dir", type=str)
+    args = parser.parse_args()
+    evaluate_abx(args)

encode.py

@@ -0,0 +1,51 @@
+import argparse
+from pathlib import Path
+import json
+import numpy as np
+import torch
+from model import Encoder
+from tqdm import tqdm
+
+
+def encode_dataset(args, params):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    model = Encoder(in_channels=params["preprocessing"]["num_mels"],
+                    encoder_channels=params["model"]["encoder_channels"],
+                    z_dim=params["model"]["z_dim"],
+                    c_dim=params["model"]["c_dim"])
+    model.to(device)
+
+    print("Load checkpoint from: {}:".format(args.checkpoint))
+    checkpoint = torch.load(args.checkpoint, map_location=lambda storage, loc: storage)
+    model.load_state_dict(checkpoint["model"])
+    model.eval()
+
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(exist_ok=True, parents=True)
+
+    hop_length_seconds = params["preprocessing"]["hop_length"] / params["preprocessing"]["sample_rate"]
+
+    in_dir = Path(args.in_dir)
+    for path in tqdm(in_dir.rglob("*.mel.npy")):
+        mel = torch.from_numpy(np.load(path)).unsqueeze(0).to(device)
+        with torch.no_grad():
+            z, c, _, _ = model(mel)
+
+        output = z.squeeze().cpu().numpy()
+        time = np.linspace(0, (mel.size(-1) - 1) * hop_length_seconds, len(output))
+        relative_path = path.relative_to(in_dir).with_suffix("")
+        out_path = out_dir / relative_path
+        out_path.parent.mkdir(exist_ok=True, parents=True)
+        np.savez(out_path.with_suffix(".npz"), features=output, time=time)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--checkpoint", type=str, help="Checkpoint path to resume")
+    parser.add_argument("--in-dir", type=str, help="Directory to encode")
+    parser.add_argument("--out-dir", type=str, help="Output path")
+    args = parser.parse_args()
+    with open("config.json") as file:
+        params = json.load(file)
+    encode_dataset(args, params)