File size: 3,325 Bytes
7339046 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | #!/usr/bin/env python3
"""Export PCA parameters from trained kicks VAE for use in HF Space.
Run this from the kicks repo root:
python export_pca.py
Outputs: pca_params.npz
"""
import numpy as np
import torch
from sklearn.decomposition import PCA
from kicks import KickDataset, KickDataloader, VAE
from kicks.cluster import extract_latents, compute_descriptors
from kicks.config import load_vae_from_checkpoint, DATA_DIR, BEST_CHECKPOINT
N_PCS = 5
DESC_KEYS = ["sub", "punch", "click", "bright", "decay"]
DESC_NAME_MAP = {
"sub": "Sub", "punch": "Punch", "click": "Click",
"bright": "Bright", "decay": "Decay",
}
def main():
device = torch.device("cpu")
print(f"Loading dataset from {DATA_DIR}...")
dataset = KickDataset(DATA_DIR)
dataloader = KickDataloader(dataset, batch_size=32, shuffle=False)
print(f"Loading VAE from {BEST_CHECKPOINT}...")
model, _ = load_vae_from_checkpoint(BEST_CHECKPOINT, device)
print("Extracting latents and spectrograms...")
latents, spectrograms = extract_latents(model, dataloader, device)
print(f" {len(latents)} samples, latent dim {latents.shape[1]}")
print("Computing perceptual descriptors for PC naming...")
descriptors = [compute_descriptors(s) for s in spectrograms]
desc_arrays = {k: np.array([d[k] for d in descriptors]) for k in DESC_KEYS}
print(f"Fitting PCA ({latents.shape[1]}D -> {N_PCS}D)...")
pca = PCA(n_components=N_PCS)
pc_projected = pca.fit_transform(latents)
# Auto-name PCs from highest descriptor correlations, flip negative axes
pc_names = []
used = set()
for i in range(N_PCS):
pc_vals = pc_projected[:, i]
best_desc, best_corr = None, 0.0
for dk in DESC_KEYS:
if dk in used:
continue
dv = desc_arrays[dk]
if pc_vals.std() > 0 and dv.std() > 0:
r = float(np.corrcoef(pc_vals, dv)[0, 1])
else:
r = 0.0
if abs(r) > abs(best_corr):
best_corr = r
best_desc = dk
if best_desc and abs(best_corr) >= 0.15:
used.add(best_desc)
pc_names.append(DESC_NAME_MAP[best_desc])
if best_corr < 0:
pca.components_[i] *= -1
pc_projected[:, i] *= -1
print(f" PC{i+1} -> {pc_names[-1]} (r={best_corr:.2f}, flipped)")
else:
print(f" PC{i+1} -> {pc_names[-1]} (r={best_corr:.2f})")
else:
pc_names.append(f"PC{i+1}")
print(f" PC{i+1} -> PC{i+1} (no strong correlation)")
pc_mins = [float(np.percentile(pc_projected[:, i], 2)) for i in range(N_PCS)]
pc_maxs = [float(np.percentile(pc_projected[:, i], 98)) for i in range(N_PCS)]
# Save
np.savez(
"pca_params.npz",
components=pca.components_,
mean=pca.mean_,
explained_variance_ratio_=pca.explained_variance_ratio_,
pc_names=np.array(pc_names, dtype=object),
pc_mins=np.array(pc_mins),
pc_maxs=np.array(pc_maxs),
)
print(f"\nSaved pca_params.npz")
print(f"PC names: {pc_names}")
print(f"PC ranges: {list(zip(pc_mins, pc_maxs))}")
print(f"Explained variance: {pca.explained_variance_ratio_}")
if __name__ == "__main__":
main()
|