/ proof-of-concept / steelsim31
Sign in
proof-of-concept / steelsim31
steelsim31 / ml / predict.py
66 lines · 2.8 KBCodeBlameHistory
7b92c9cAdd dataset generation scripts, ML fitting code and reference recordingsJeremy Magland 1"""Applies a trained model to recordings.
3 python3 ml/predict.py --model RUN_DIR/model.pt [--logf] REAL.wav [more.wav ...]
5Each recording is converted to mono at 48 kHz, cut to 3 s from its onset (see
6match.py), peak-normalized, turned into the spectrogram the model was trained
7on, and passed through the network. To show the sensitivity to the alignment,
8the prediction is repeated with the onset shifted by a few milliseconds and
9the spread is reported alongside the value at the detected onset.
10"""
12import argparse
13import math
14import os
15import sys
17import numpy as np
18import torch
20sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
21from features import spectrogram, spectrogram_logf # noqa: E402
22from match import read_wav_any, onset, segment, FS # noqa: E402
23from train import Net, SpecNet, TARGETS, UNITS # noqa: E402
25torch.backends.cudnn.enabled = False
28def main():
29 ap = argparse.ArgumentParser()
30 ap.add_argument("--model", required=True)
31 ap.add_argument("--logf", action="store_true")
32 ap.add_argument("--shifts", default="0,-5,5,10,20", help="onset shifts in ms for the sensitivity check")
33 ap.add_argument("wavs", nargs="+")
34 args = ap.parse_args()
35 ck = torch.load(args.model, map_location="cpu", weights_only=False)
36 net = SpecNet(len(ck["targets"]), n_freq=ck["n_freq"]) if ck.get("arch") == "spec1d" else Net(len(ck["targets"]))
37 net.load_state_dict(ck["state"])
38 net.eval()
39 mu, sd = ck["mu"], ck["sd"]
40 Tinv = np.linalg.inv(ck["T"]) if "T" in ck else np.eye(len(ck["targets"]))
41 logf = ck["input"] == "logf" if "input" in ck else args.logf
42 spec = spectrogram_logf if logf else spectrogram
43 shifts = [float(s) for s in args.shifts.split(",")]
44 for path in args.wavs:
45 x = read_wav_any(path)
46 st = onset(x)
47 preds = []
48 for sh in shifts:
49 s0 = max(0, st + int(sh * FS / 1000))
50 X = torch.from_numpy(spec(segment(x, s0)).astype(np.float32))[None, None]
51 with torch.no_grad():
52 z = net(X)[0].numpy()
53 preds.append((z * sd + mu) @ Tinv.T)
54 preds = np.array(preds)
55 p0, lo, hi = preds[0], preds.min(0), preds.max(0)
56 print(f"{path}: {len(x) / FS:.2f} s, onset at {st / FS:.3f} s")
57 for i, t in enumerate(TARGETS):
58 unit, scale = UNITS[t]
59 print(f" {t:>4}: {math.exp(p0[i]) * scale:8.2f} {unit:<7} (onset shifts give {math.exp(lo[i]) * scale:.2f} to {math.exp(hi[i]) * scale:.2f})")
60 d = {t: p0[i] for i, t in enumerate(TARGETS)}
61 kappa = math.exp(d["h"] + 0.5 * (d["E"] - d["rho"])) / math.sqrt(12 * (1 - 0.29 ** 2))
62 print(f" E/rho {math.exp(d['E'] - d['rho']) / 1e6:.1f} MPa m^3/kg, kappa {kappa:.3f} m^2/s, area {math.exp(d['Lx'] + d['Ly']) / 0.0254 ** 2:.0f} in^2, aspect {math.exp(d['Lx'] - d['Ly']):.2f}")
65if __name__ == "__main__":
66 main()
↑↓ move↵openescclose