478 lines
16 KiB
Python
478 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
"""Python verifier for Tom game's tiny keyword spotting model.
|
|
|
|
This script mirrors src/recognition/TinyKwsRecognizer.cpp closely enough for
|
|
Windows-side validation without building TensorFlow Lite C++ on Windows.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import math
|
|
import sys
|
|
import wave
|
|
from pathlib import Path
|
|
from typing import Iterable, Tuple
|
|
|
|
import numpy as np
|
|
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[4]
|
|
DEFAULT_MODEL = (
|
|
REPO_ROOT
|
|
/ "src"
|
|
/ "Apps"
|
|
/ "Game"
|
|
/ "model"
|
|
/ "mlcommons-tiny-kws"
|
|
/ "models"
|
|
/ "kws_ref_model_float32.tflite"
|
|
)
|
|
|
|
LABELS = [
|
|
"Down",
|
|
"Go",
|
|
"Left",
|
|
"No",
|
|
"Off",
|
|
"On",
|
|
"Right",
|
|
"Stop",
|
|
"Up",
|
|
"Yes",
|
|
"Silence",
|
|
"Unknown",
|
|
]
|
|
|
|
COMMAND_LABELS = {"Down", "Go", "Left", "No", "On", "Right", "Stop", "Up", "Yes"}
|
|
|
|
TARGET_SAMPLE_RATE = 16000
|
|
TARGET_CHANNELS = 1
|
|
WINDOW_SAMPLES = 16000
|
|
WINDOW_STEP_SAMPLES = 8000
|
|
FFT_SIZE = 512
|
|
SPECTRUM_BINS = FFT_SIZE // 2 + 1
|
|
FRAME_STEP_SAMPLES = 320
|
|
FEATURE_FRAME_COUNT = 49
|
|
MEL_BIN_COUNT = 40
|
|
MFCC_BIN_COUNT = 10
|
|
|
|
|
|
def load_interpreter(model_path: Path):
|
|
try:
|
|
from tflite_runtime.interpreter import Interpreter
|
|
|
|
backend = "tflite_runtime"
|
|
except ImportError:
|
|
try:
|
|
import tensorflow as tf
|
|
|
|
Interpreter = tf.lite.Interpreter
|
|
backend = "tensorflow.lite"
|
|
except Exception as exc:
|
|
raise RuntimeError(
|
|
"No TFLite interpreter backend is available. Install either "
|
|
"`tflite-runtime` or `tensorflow` in the active Conda env."
|
|
) from exc
|
|
|
|
interpreter = Interpreter(model_path=str(model_path), num_threads=1)
|
|
interpreter.allocate_tensors()
|
|
return backend, interpreter
|
|
|
|
|
|
def read_wav_int16(path: Path) -> Tuple[np.ndarray, int, int]:
|
|
with wave.open(str(path), "rb") as wav:
|
|
sample_width = wav.getsampwidth()
|
|
if sample_width != 2:
|
|
raise ValueError(f"{path} is {sample_width * 8}-bit PCM; expected 16-bit PCM")
|
|
|
|
sample_rate = wav.getframerate()
|
|
channels = wav.getnchannels()
|
|
frame_count = wav.getnframes()
|
|
raw = wav.readframes(frame_count)
|
|
|
|
samples = np.frombuffer(raw, dtype="<i2").astype(np.int16, copy=False)
|
|
return samples, sample_rate, channels
|
|
|
|
|
|
def write_wav_int16(path: Path, samples: np.ndarray, sample_rate: int, channels: int) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
with wave.open(str(path), "wb") as wav:
|
|
wav.setnchannels(channels)
|
|
wav.setsampwidth(2)
|
|
wav.setframerate(sample_rate)
|
|
wav.writeframes(np.asarray(samples, dtype="<i2").tobytes())
|
|
|
|
|
|
def list_devices() -> None:
|
|
try:
|
|
import sounddevice as sd
|
|
except ImportError as exc:
|
|
raise RuntimeError("`sounddevice` is required for --list-devices") from exc
|
|
|
|
print(sd.query_devices())
|
|
|
|
|
|
def record_audio(seconds: float, sample_rate: int, channels: int, device: str | None) -> np.ndarray:
|
|
try:
|
|
import sounddevice as sd
|
|
except ImportError as exc:
|
|
raise RuntimeError("`sounddevice` is required for --record") from exc
|
|
|
|
frame_count = int(round(seconds * sample_rate))
|
|
print(
|
|
f"[INFO] Recording {seconds:.2f}s at {sample_rate} Hz, "
|
|
f"{channels} channel(s), device={device if device is not None else 'default'}"
|
|
)
|
|
audio = sd.rec(
|
|
frame_count,
|
|
samplerate=sample_rate,
|
|
channels=channels,
|
|
dtype="int16",
|
|
device=device,
|
|
)
|
|
sd.wait()
|
|
return np.asarray(audio, dtype=np.int16).reshape(-1)
|
|
|
|
|
|
def convert_channels(samples: np.ndarray, source_channels: int, target_channels: int) -> np.ndarray:
|
|
if samples.size == 0 or source_channels == target_channels:
|
|
return samples.astype(np.int16, copy=False)
|
|
|
|
frames = samples.reshape((-1, source_channels))
|
|
if target_channels == 1:
|
|
mixed = np.mean(frames.astype(np.int32), axis=1)
|
|
return np.clip(np.round(mixed), -32768, 32767).astype(np.int16)
|
|
|
|
output = np.zeros((frames.shape[0], target_channels), dtype=np.int16)
|
|
for channel in range(target_channels):
|
|
output[:, channel] = frames[:, min(channel, source_channels - 1)]
|
|
return output.reshape(-1)
|
|
|
|
|
|
def resample_linear(
|
|
samples: np.ndarray,
|
|
source_rate: int,
|
|
target_rate: int,
|
|
channels: int,
|
|
) -> np.ndarray:
|
|
if samples.size == 0 or source_rate == target_rate:
|
|
return samples.astype(np.int16, copy=False)
|
|
|
|
frames = samples.reshape((-1, channels)).astype(np.float32)
|
|
source_frame_count = frames.shape[0]
|
|
target_frame_count = max(1, int(math.ceil(source_frame_count * target_rate / source_rate)))
|
|
|
|
source_positions = np.arange(target_frame_count, dtype=np.float32) * (source_rate / target_rate)
|
|
source0 = np.minimum(source_positions.astype(np.int64), source_frame_count - 1)
|
|
source1 = np.minimum(source0 + 1, source_frame_count - 1)
|
|
t = (source_positions - source0).reshape((-1, 1))
|
|
|
|
mixed = frames[source0] + (frames[source1] - frames[source0]) * t
|
|
return np.clip(mixed, -32768, 32767).astype(np.int16).reshape(-1)
|
|
|
|
|
|
def trim_silence(samples: np.ndarray, threshold: float, channels: int) -> np.ndarray:
|
|
if samples.size == 0:
|
|
return samples
|
|
|
|
frames = samples.reshape((-1, channels))
|
|
loud = np.any(np.abs(frames.astype(np.float32) / 32768.0) >= threshold, axis=1)
|
|
indices = np.flatnonzero(loud)
|
|
if indices.size == 0:
|
|
return samples
|
|
|
|
start = int(indices[0])
|
|
end = int(indices[-1]) + 1
|
|
return frames[start:end].reshape(-1).astype(np.int16, copy=False)
|
|
|
|
|
|
def prepare_audio(samples: np.ndarray, sample_rate: int, channels: int, trim: bool = True) -> np.ndarray:
|
|
if samples.size == 0 or sample_rate <= 0 or channels <= 0:
|
|
return np.array([], dtype=np.int16)
|
|
|
|
prepared = samples.astype(np.int16, copy=False)
|
|
prepared_channels = channels
|
|
|
|
if prepared_channels != TARGET_CHANNELS:
|
|
prepared = convert_channels(prepared, prepared_channels, TARGET_CHANNELS)
|
|
prepared_channels = TARGET_CHANNELS
|
|
|
|
if sample_rate != TARGET_SAMPLE_RATE:
|
|
prepared = resample_linear(prepared, sample_rate, TARGET_SAMPLE_RATE, prepared_channels)
|
|
|
|
if not trim:
|
|
return prepared
|
|
|
|
trimmed = trim_silence(prepared, 0.02, TARGET_CHANNELS)
|
|
return trimmed if trimmed.size > 0 else prepared
|
|
|
|
|
|
def hertz_to_mel(hertz: np.ndarray | float) -> np.ndarray | float:
|
|
return 2595.0 * np.log10(1.0 + np.asarray(hertz) / 700.0)
|
|
|
|
|
|
def mel_to_hertz(mel: np.ndarray | float) -> np.ndarray | float:
|
|
return 700.0 * (np.power(10.0, np.asarray(mel) / 2595.0) - 1.0)
|
|
|
|
|
|
def make_mel_filter_bank() -> np.ndarray:
|
|
filters = np.zeros((MEL_BIN_COUNT, SPECTRUM_BINS), dtype=np.float32)
|
|
min_mel = float(hertz_to_mel(20.0))
|
|
max_mel = float(hertz_to_mel(4000.0))
|
|
mel_points = np.linspace(min_mel, max_mel, MEL_BIN_COUNT + 2, dtype=np.float32)
|
|
hertz_points = mel_to_hertz(mel_points).astype(np.float32)
|
|
frequencies = (
|
|
np.arange(SPECTRUM_BINS, dtype=np.float32) * TARGET_SAMPLE_RATE / float(FFT_SIZE)
|
|
)
|
|
|
|
for mel in range(MEL_BIN_COUNT):
|
|
left = hertz_points[mel]
|
|
center = hertz_points[mel + 1]
|
|
right = hertz_points[mel + 2]
|
|
|
|
if center > left:
|
|
rising = (frequencies >= left) & (frequencies <= center)
|
|
filters[mel, rising] = (frequencies[rising] - left) / (center - left)
|
|
|
|
if right > center:
|
|
falling = (frequencies > center) & (frequencies <= right)
|
|
filters[mel, falling] = (right - frequencies[falling]) / (right - center)
|
|
|
|
return filters
|
|
|
|
|
|
HANN_WINDOW = (0.5 - 0.5 * np.cos((2.0 * np.pi * np.arange(FFT_SIZE)) / (FFT_SIZE - 1))).astype(
|
|
np.float32
|
|
)
|
|
MEL_FILTER_BANK = make_mel_filter_bank()
|
|
DCT_BASIS = np.zeros((MFCC_BIN_COUNT, MEL_BIN_COUNT), dtype=np.float32)
|
|
for coeff in range(MFCC_BIN_COUNT):
|
|
norm = math.sqrt(1.0 / MEL_BIN_COUNT) if coeff == 0 else math.sqrt(2.0 / MEL_BIN_COUNT)
|
|
for mel in range(MEL_BIN_COUNT):
|
|
DCT_BASIS[coeff, mel] = norm * math.cos(
|
|
math.pi * (mel + 0.5) * coeff / float(MEL_BIN_COUNT)
|
|
)
|
|
|
|
|
|
def normalize_window(samples: np.ndarray, start: int) -> np.ndarray:
|
|
output = np.zeros(WINDOW_SAMPLES, dtype=np.float32)
|
|
available = max(0, min(WINDOW_SAMPLES, samples.size - start))
|
|
if available <= 0:
|
|
return output
|
|
|
|
window = samples[start : start + available].astype(np.float32)
|
|
max_abs = float(np.max(np.abs(window)))
|
|
if max_abs <= 0.0:
|
|
return output
|
|
|
|
output[:available] = window / max_abs
|
|
return output
|
|
|
|
|
|
def extract_mfcc_features(window_samples: np.ndarray) -> np.ndarray:
|
|
features = np.zeros((FEATURE_FRAME_COUNT, MFCC_BIN_COUNT), dtype=np.float32)
|
|
|
|
for frame in range(FEATURE_FRAME_COUNT):
|
|
offset = frame * FRAME_STEP_SAMPLES
|
|
frame_samples = np.zeros(FFT_SIZE, dtype=np.float32)
|
|
available = max(0, min(FFT_SIZE, window_samples.size - offset))
|
|
if available > 0:
|
|
frame_samples[:available] = window_samples[offset : offset + available]
|
|
|
|
spectrum = np.abs(np.fft.rfft(frame_samples * HANN_WINDOW, n=FFT_SIZE)).astype(np.float32)
|
|
mel_energy = MEL_FILTER_BANK @ spectrum
|
|
log_mel = np.log(mel_energy + 1.0e-6).astype(np.float32)
|
|
features[frame, :] = DCT_BASIS @ log_mel
|
|
|
|
return features
|
|
|
|
|
|
def tensor_from_features(features: np.ndarray, input_details: dict) -> np.ndarray:
|
|
input_shape = tuple(int(v) for v in input_details["shape"])
|
|
tensor = features.astype(np.float32).reshape(input_shape)
|
|
dtype = input_details["dtype"]
|
|
|
|
if dtype == np.float32:
|
|
return tensor
|
|
|
|
scale, zero_point = input_details.get("quantization", (0.0, 0))
|
|
if scale == 0:
|
|
raise ValueError(f"Unsupported quantized input without scale: {input_details}")
|
|
|
|
quantized = np.round(tensor / scale + zero_point)
|
|
info = np.iinfo(dtype)
|
|
return np.clip(quantized, info.min, info.max).astype(dtype)
|
|
|
|
|
|
def dequantize_output(output: np.ndarray, output_details: dict) -> np.ndarray:
|
|
if output.dtype == np.float32:
|
|
return output.astype(np.float32).reshape(-1)
|
|
|
|
scale, zero_point = output_details.get("quantization", (0.0, 0))
|
|
if scale == 0:
|
|
return output.astype(np.float32).reshape(-1)
|
|
|
|
return ((output.astype(np.float32) - zero_point) * scale).reshape(-1)
|
|
|
|
|
|
def run_window(interpreter, window_samples: np.ndarray) -> np.ndarray:
|
|
input_details = interpreter.get_input_details()[0]
|
|
output_details = interpreter.get_output_details()[0]
|
|
features = extract_mfcc_features(window_samples)
|
|
tensor = tensor_from_features(features, input_details)
|
|
|
|
interpreter.set_tensor(input_details["index"], tensor)
|
|
interpreter.invoke()
|
|
raw_output = interpreter.get_tensor(output_details["index"])
|
|
return dequantize_output(raw_output, output_details)
|
|
|
|
|
|
def window_starts(sample_count: int) -> Iterable[int]:
|
|
if sample_count <= WINDOW_SAMPLES:
|
|
yield 0
|
|
return
|
|
|
|
start = 0
|
|
while start + WINDOW_SAMPLES <= sample_count:
|
|
yield start
|
|
start += WINDOW_STEP_SAMPLES
|
|
|
|
last_start = sample_count - WINDOW_SAMPLES
|
|
if last_start % WINDOW_STEP_SAMPLES != 0:
|
|
yield last_start
|
|
|
|
|
|
def recognize(
|
|
interpreter,
|
|
samples: np.ndarray,
|
|
sample_rate: int,
|
|
channels: int,
|
|
threshold: float,
|
|
trim: bool,
|
|
) -> Tuple[str, float, np.ndarray, int]:
|
|
prepared = prepare_audio(samples, sample_rate, channels, trim=trim)
|
|
if prepared.size == 0:
|
|
raise ValueError("No audio samples available after preprocessing")
|
|
|
|
best_command_label = "None"
|
|
best_command_confidence = 0.0
|
|
best_command_probs = np.zeros(len(LABELS), dtype=np.float32)
|
|
best_command_start = 0
|
|
|
|
# 用于显示：取所有滑窗中“最高类置信度”最大的那一窗，
|
|
# 这样即使最高类不是游戏命令，print_topk 也不会显示全 0。
|
|
best_overall_probs = np.zeros(len(LABELS), dtype=np.float32)
|
|
best_overall_max_confidence = 0.0
|
|
|
|
for start in window_starts(prepared.size):
|
|
probs = run_window(interpreter, normalize_window(prepared, start))
|
|
if probs.size < len(LABELS):
|
|
raise ValueError(f"Model returned {probs.size} values; expected at least {len(LABELS)}")
|
|
|
|
probs_slice = probs[: len(LABELS)]
|
|
class_index = int(np.argmax(probs_slice))
|
|
label = LABELS[class_index]
|
|
confidence = float(probs_slice[class_index])
|
|
is_command = label in COMMAND_LABELS
|
|
|
|
if confidence > best_overall_max_confidence:
|
|
best_overall_max_confidence = confidence
|
|
best_overall_probs = probs_slice.copy()
|
|
|
|
if is_command and confidence > best_command_confidence:
|
|
best_command_label = label
|
|
best_command_confidence = confidence
|
|
best_command_probs = probs_slice.copy()
|
|
best_command_start = start
|
|
|
|
if best_command_confidence < threshold:
|
|
return "None", best_command_confidence, best_overall_probs, best_command_start
|
|
|
|
return best_command_label, best_command_confidence, best_overall_probs, best_command_start
|
|
|
|
|
|
def print_topk(probs: np.ndarray, topk: int) -> None:
|
|
top_indices = np.argsort(probs[: len(LABELS)])[::-1][:topk]
|
|
print("[INFO] Top classes:")
|
|
for index in top_indices:
|
|
print(f" {LABELS[int(index)]:>7s}: {float(probs[int(index)]):.6f}")
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(
|
|
description="Run Tom game's tiny KWS TFLite model through Python."
|
|
)
|
|
parser.add_argument("--model", type=Path, default=DEFAULT_MODEL, help="Path to .tflite model")
|
|
parser.add_argument("--wav", type=Path, help="16-bit PCM WAV input")
|
|
parser.add_argument("--record", type=float, help="Record N seconds from microphone")
|
|
parser.add_argument("--device", help="sounddevice input device index/name for --record")
|
|
parser.add_argument("--mic-rate", type=int, default=48000, help="Microphone sample rate")
|
|
parser.add_argument("--channels", type=int, default=1, help="Input channel count")
|
|
parser.add_argument("--save-wav", type=Path, help="Save recorded raw microphone audio")
|
|
parser.add_argument("--threshold", type=float, default=0.75, help="Command confidence threshold")
|
|
parser.add_argument("--topk", type=int, default=5, help="Number of classes to print")
|
|
parser.add_argument("--no-trim", action="store_true", help="Disable silence trimming before KWS")
|
|
parser.add_argument("--list-devices", action="store_true", help="List sounddevice devices and exit")
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
|
|
try:
|
|
if args.list_devices:
|
|
list_devices()
|
|
return 0
|
|
|
|
if args.wav is None and args.record is None:
|
|
raise ValueError("Specify either --wav or --record, or use --list-devices")
|
|
|
|
if args.wav is not None and args.record is not None:
|
|
raise ValueError("Use only one input mode: --wav or --record")
|
|
|
|
model_path = args.model.resolve()
|
|
if not model_path.exists():
|
|
raise FileNotFoundError(f"Model not found: {model_path}")
|
|
|
|
backend, interpreter = load_interpreter(model_path)
|
|
print(f"[INFO] Interpreter backend: {backend}")
|
|
print(f"[INFO] Model: {model_path}")
|
|
print(f"[INFO] Input details: {interpreter.get_input_details()[0]}")
|
|
print(f"[INFO] Output details: {interpreter.get_output_details()[0]}")
|
|
|
|
if args.wav is not None:
|
|
samples, sample_rate, channels = read_wav_int16(args.wav.resolve())
|
|
print(f"[INFO] WAV: {args.wav.resolve()} ({sample_rate} Hz, {channels} channel(s))")
|
|
else:
|
|
sample_rate = args.mic_rate
|
|
channels = args.channels
|
|
samples = record_audio(args.record, sample_rate, channels, args.device)
|
|
if args.save_wav is not None:
|
|
write_wav_int16(args.save_wav.resolve(), samples, sample_rate, channels)
|
|
print(f"[INFO] Saved recording: {args.save_wav.resolve()}")
|
|
|
|
label, confidence, probs, start = recognize(
|
|
interpreter,
|
|
samples,
|
|
sample_rate,
|
|
channels,
|
|
args.threshold,
|
|
trim=not args.no_trim,
|
|
)
|
|
|
|
print_topk(probs, max(1, args.topk))
|
|
print(
|
|
f"[RESULT] command={label} confidence={confidence:.6f} "
|
|
f"window_start_sample={start} threshold={args.threshold:.3f}"
|
|
)
|
|
return 0
|
|
|
|
except Exception as exc:
|
|
print(f"[ERROR] {exc}", file=sys.stderr)
|
|
return 2
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|