File size: 2,030 Bytes
37c4768 | 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 | """Example: Run DiariZen segmentation on a WAV file."""
import argparse
import sys
import numpy as np
from diarizen_sdk import DiarizenSegmenter
from diarizen_sdk.postprocess import log_probs_to_probs, top_speakers_at_frame
def main():
parser = argparse.ArgumentParser(description="DiariZen speaker segmentation")
parser.add_argument("audio", help="Path to 16kHz mono WAV file")
parser.add_argument("--cnn-model", default="cnn_features.axmodel",
help="Path to CNN NPU model")
parser.add_argument("--backend-model", default="backend.onnx",
help="Path to backend ONNX model")
args = parser.parse_args()
# Load audio
try:
import soundfile as sf
audio, sr = sf.read(args.audio, dtype="float32")
except ImportError:
print("soundfile not available, using scipy.io.wavfile")
from scipy.io import wavfile
sr, audio = wavfile.read(args.audio)
audio = audio.astype(np.float32) / 32768.0
if audio.ndim > 1:
audio = audio[:, 0] # Use first channel
print(f"Audio: {len(audio)} samples @ {sr} Hz")
# Create segmenter and run
segmenter = DiarizenSegmenter(args.cnn_model, args.backend_model)
log_probs = segmenter(audio, sr)
print(f"Output shape: {log_probs.shape}")
print(f" Frames: {log_probs.shape[1]}, Classes: {log_probs.shape[2]}")
# Show results for key frames
probs = log_probs_to_probs(log_probs)
check_frames = [0, 50, 100, 150, 198]
print("\nTop-3 speaker classes per selected frame:")
for f in check_frames:
top = top_speakers_at_frame(log_probs, f, top_k=3)
items = ", ".join(f"cls {c}: {lp:.2f}" for c, lp in top)
print(f" Frame {f:3d}: {items}")
# Overall most active class
mean_probs = probs[0].mean(axis=0)
top_class = int(np.argmax(mean_probs))
print(f"\nMost active class overall: {top_class} "
f"(avg prob={mean_probs[top_class]:.4f})")
if __name__ == "__main__":
main()
|