File size: 5,498 Bytes
001425d | 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 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | #!/usr/bin/env python3
"""
Test script to run and compare both PyTorch and LiteRT / TFLite inference.
Synthesizes a realistic audio chirp sample if no test WAV is provided.
"""
import argparse
import json
from pathlib import Path
import numpy as np
import scipy.io.wavfile as wavfile
import tensorflow as tf
from src.inference import predict_wav, identify, _summed, genus_of, group_of, DEFAULT_THRESHOLD
from src.audio import read_audio, audio_quality, resample, crop_offsets, logmel, SR
ROOT = Path(__file__).resolve().parent
def generate_synthetic_chirp(path, duration=3.0, sr=16000):
"""Generate a 3-second synthetic insect chirp WAV for local test inference."""
t = np.linspace(0, duration, int(sr * duration), endpoint=False)
# Carrier frequency around 4500 Hz modulated with bursts (simulating cricket/cicada pulse train)
carrier = np.sin(2 * np.pi * 4500 * t)
pulse = (np.sin(2 * np.pi * 25 * t) > 0.3).astype(np.float32)
waveform = (carrier * pulse * 0.4).astype(np.float32)
# Save as 16-bit PCM WAV
wavfile.write(str(path), sr, (waveform * 32767).astype(np.int16))
return path
def predict_tflite(path, model_dir=ROOT, tflite_path=None):
"""Run inference using the LiteRT / TFLite model."""
model_dir = Path(model_dir)
cfg = json.loads((model_dir / 'config.json').read_text())
labels = json.loads((model_dir / 'labels.json').read_text())
if tflite_path is None:
tflite_path = model_dir / 'model.tflite'
else:
tflite_path = Path(tflite_path)
if not tflite_path.is_file():
raise FileNotFoundError(f"LiteRT model not found at '{tflite_path}'. Run export_litert.py first.")
interpreter = tf.lite.Interpreter(model_path=str(tflite_path))
interpreter.allocate_tensors()
in_idx = interpreter.get_input_details()[0]['index']
out_idx = interpreter.get_output_details()[0]['index']
sr, y = read_audio(path)
ok, reasons = audio_quality(y)
y16 = resample(y, sr)
offs = crop_offsets(len(y16), int(cfg.get('eval_crops', 1)))
logits_list = []
for off in offs:
feat = logmel(y16, SR, off) # (1, 64, 201)
feat_batch = feat[None, :, :, :].astype(np.float32) # (1, 1, 64, 201)
interpreter.set_tensor(in_idx, feat_batch)
interpreter.invoke()
logits = interpreter.get_tensor(out_idx) # (1, 32)
logits_list.append(logits[0])
# Softmax over logits, then mean across crops
logits_arr = np.array(logits_list)
exp_logits = np.exp(logits_arr - np.max(logits_arr, axis=-1, keepdims=True))
crop_probs = exp_logits / np.sum(exp_logits, axis=-1, keepdims=True)
probs = crop_probs.mean(axis=0)
order = np.argsort(probs)[::-1][:5]
thr = float(cfg.get('confidence_threshold', DEFAULT_THRESHOLD))
return {
'identification': identify(probs, labels, thr),
'genus_scores': [{'genus': g, 'score': v} for g, v in _summed(probs, labels, genus_of)[:3]],
'group_scores': [{'group': g, 'score': v} for g, v in _summed(probs, labels, group_of)],
'predicted_species': labels[str(int(order[0]))],
'confidence': float(probs[order[0]]),
'top_k': [{'species': labels[str(int(i))], 'score': float(probs[i])} for i in order],
'quality_pass': bool(ok),
'quality_reasons': reasons,
'model_version': cfg['model_version'],
'runtime': 'LiteRT / TFLite'
}
def main():
parser = argparse.ArgumentParser(description="Test PyTorch and LiteRT insect classification inference")
parser.add_argument('--wav', default=None, help="Path to input audio file (.wav). If not provided, a sample chirp is created.")
parser.add_argument('--model-dir', default='.', help="Model directory")
args = parser.parse_args()
model_dir = Path(args.model_dir)
wav_path = args.wav
if wav_path is None:
wav_path = model_dir / 'test_sample.wav'
print(f"[*] Generating synthetic 3-second insect test chirp at '{wav_path}'...")
generate_synthetic_chirp(wav_path)
else:
wav_path = Path(wav_path)
print(f"\n[1] Running PyTorch Inference on '{wav_path}'...")
pt_res = predict_wav(wav_path, model_dir)
print(f" PyTorch Top-1 Species : {pt_res['predicted_species']}")
print(f" PyTorch Confidence : {pt_res['confidence']:.4f}")
print(f" PyTorch Identification: {pt_res['identification']}")
print(f"\n[2] Running LiteRT / TFLite Inference on '{wav_path}'...")
tflite_res = predict_tflite(wav_path, model_dir)
print(f" LiteRT Top-1 Species : {tflite_res['predicted_species']}")
print(f" LiteRT Confidence : {tflite_res['confidence']:.4f}")
print(f" LiteRT Identification : {tflite_res['identification']}")
# Parity check
species_match = pt_res['predicted_species'] == tflite_res['predicted_species']
conf_diff = abs(pt_res['confidence'] - tflite_res['confidence'])
print("\n" + "=" * 60)
print(" VERIFICATION SUMMARY:")
print(f" - Top Predicted Species Match : {'YES' if species_match else 'NO'} ({pt_res['predicted_species']})")
print(f" - Confidence Difference : {conf_diff:.8e}")
print(f" - Audio Quality Check Passed : {pt_res['quality_pass']}")
print("=" * 60)
if species_match and conf_diff < 1e-4:
print("[+] Test PASSED: Both runtimes agree with exact parity.\n")
else:
print("[-] Test WARNING: Slight divergence detected.\n")
if __name__ == '__main__':
main()
|