test_server.py

#!/usr/bin/env python3
"""
VoiceSwitch Test Server
========================
Minimal HTTP API server for testing the KWS model from a browser.
Loads the trained model and serves predictions via REST endpoint.

Usage:
    python test_server.py              # Start on port 8080
    python test_server.py --port 9000  # Custom port
    python test_server.py --cors       # Enable CORS for all origins

Endpoints:
    POST /predict     - Send audio, get prediction back
    GET  /health      - Server health check
    GET  /            - Serves the test UI page
"""

import sys
import json
import wave
import io
import argparse
from pathlib import Path

import numpy as np

# Ensure project root in path
sys.path.insert(0, str(Path(__file__).resolve().parent))

import os
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

import tensorflow as tf
import yaml

from edge.inference import EdgeMFCC

# === Configuration ===
MODEL_PATH = "models/dscnn_checkpoint.keras"  # Bilingual Yoruba+English model
CONFIG_PATH = "edge/config.yaml"
SAMPLE_RATE = 16000
WINDOW_SAMPLES = 16000  # 1 second


class VoiceSwitchPredictor:
    """Loads model and MFCC extractor, runs inference on audio."""

    def __init__(self, model_path=MODEL_PATH, config_path=CONFIG_PATH):
        print(f"Loading model: {model_path}")
        self.model = tf.keras.models.load_model(model_path)
        self.model.trainable = False
        print(f"  Params: {self.model.count_params():,}")

        with open(config_path) as f:
            self.config = yaml.safe_load(f)
        self.labels = self.config["model"]["labels"]
        self.mfcc = EdgeMFCC(self.config)

        # Warm up with dummy data
        dummy = np.random.randn(WINDOW_SAMPLES).astype(np.float32) * 0.001
        self.predict(dummy)
        print("  Model loaded and ready.")

    def predict(self, audio):
        """
        Run inference on raw audio samples.
        Args:
            audio: numpy array of float32 samples in [-1, 1], any length
        Returns:
            dict with label, confidence, and all probabilities
        """
        # Convert to float32
        audio = audio.astype(np.float32)

        # Normalize if int16 range
        if np.max(np.abs(audio)) > 1.5:
            audio = audio / 32767.0

        # Pad or trim to exactly 1 second
        if len(audio) < WINDOW_SAMPLES:
            audio = np.pad(audio, (0, WINDOW_SAMPLES - len(audio)))
        elif len(audio) > WINDOW_SAMPLES:
            # Take the center 1 second
            start = (len(audio) - WINDOW_SAMPLES) // 2
            audio = audio[start:start + WINDOW_SAMPLES]

        # Extract MFCC features
        features = self.mfcc.extract(audio)
        features = np.expand_dims(features, axis=0)

        # Run inference
        logits = self.model(features, training=False)
        probs = tf.nn.softmax(logits, axis=-1).numpy()[0]

        class_id = int(np.argmax(probs))
        confidence = float(probs[class_id])

        return {
            "class_id": class_id,
            "label": self.labels[class_id] if class_id < len(self.labels) else "unknown",
            "confidence": round(confidence, 4),
            "probabilities": {
                label: round(float(prob), 4)
                for label, prob in zip(self.labels, probs)
            },
            "rms_energy": round(float(np.sqrt(np.mean(audio ** 2))), 6),
            "duration_samples": len(audio),
        }


def create_app(predictor):
    """Create a minimal WSGI app using Python's built-in http.server."""
    from http.server import HTTPServer, BaseHTTPRequestHandler

    UI_HTML = (Path(__file__).resolve().parent / "test_ui.html").read_text(encoding="utf-8")

    class RequestHandler(BaseHTTPRequestHandler):
        def log_message(self, format, *args):
            # Quieter logging
            if "/predict" in str(args):
                print(f"  [API] {args[0]}")

        def do_GET(self):
            if self.path == "/" or self.path == "/index.html":
                self.send_response(200)
                self.send_header("Content-Type", "text/html; charset=utf-8")
                self.send_header("Access-Control-Allow-Origin", "*")
                self.end_headers()
                self.wfile.write(UI_HTML.encode("utf-8"))
            elif self.path == "/health":
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.send_header("Access-Control-Allow-Origin", "*")
                self.end_headers()
                resp = json.dumps({
                    "status": "ok",
                    "model": str(Path(MODEL_PATH).name),
                    "labels": predictor.labels,
                    "params": predictor.model.count_params(),
                })
                self.wfile.write(resp.encode("utf-8"))
            else:
                self.send_response(404)
                self.end_headers()

        def do_POST(self):
            if self.path == "/predict":
                try:
                    content_length = int(self.headers.get("Content-Length", 0))
                    content_type = self.headers.get("Content-Type", "")

                    if "multipart/form-data" in content_type:
                        # Parse multipart form data with audio file
                        audio_data = self._parse_multipart(content_type, self.rfile.read(content_length))
                    else:
                        # Raw WAV bytes
                        audio_data = self.rfile.read(content_length)
                        audio_data = self._decode_wav(audio_data)

                    if audio_data is None or len(audio_data) < 100:
                        self._send_error(400, "Invalid audio data")
                        return

                    result = predictor.predict(audio_data)

                    self.send_response(200)
                    self.send_header("Content-Type", "application/json")
                    self.send_header("Access-Control-Allow-Origin", "*")
                    self.end_headers()
                    self.wfile.write(json.dumps(result).encode("utf-8"))

                except Exception as e:
                    print(f"  [ERROR] {e}")
                    self._send_error(500, str(e))
            else:
                self.send_response(404)
                self.end_headers()

        def do_OPTIONS(self):
            self.send_response(200)
            self.send_header("Access-Control-Allow-Origin", "*")
            self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
            self.send_header("Access-Control-Allow-Headers", "Content-Type")
            self.end_headers()

        def _parse_multipart(self, content_type, body):
            """Simple multipart parser for audio file upload."""
            # Find boundary
            boundary = None
            for part in content_type.split(";"):
                part = part.strip()
                if part.startswith("boundary="):
                    boundary = part.split("=", 1)[1].strip('"')
                    break

            if not boundary:
                return None

            boundary_bytes = boundary.encode("utf-8")
            parts = body.split(b"--" + boundary_bytes)

            for part in parts:
                if b"Content-Disposition" in part:
                    # Find the double CRLF separating headers from body
                    header_end = part.find(b"\r\n\r\n")
                    if header_end == -1:
                        continue
                    body_start = header_end + 4
                    raw_audio = part[body_start:]

                    # Strip trailing boundary
                    end = raw_audio.rfind(b"\r\n")
                    if end > 0:
                        raw_audio = raw_audio[:end]

                    return self._decode_wav(raw_audio)

            return None

        def _decode_wav(self, data):
            """Decode WAV bytes to numpy float32 array."""
            try:
                with wave.open(io.BytesIO(data)) as w:
                    n_frames = w.getnframes()
                    n_channels = w.getnchannels()
                    sample_width = w.getsampwidth()
                    framerate = w.getframerate()

                    raw = w.readframes(n_frames)
                    fmt = {1: "b", 2: "h", 4: "i"}[sample_width]
                    audio = np.frombuffer(raw, dtype=f"<i{sample_width}")

                    if n_channels > 1:
                        audio = audio.reshape(-1, n_channels)
                        audio = audio.mean(axis=1)

                    audio = audio.astype(np.float32) / (2 ** (8 * sample_width - 1))

                    # Resample if needed
                    if framerate != SAMPLE_RATE:
                        from scipy.signal import resample_poly
                        import math
                        g = math.gcd(framerate, SAMPLE_RATE)
                        audio = resample_poly(
                            audio.astype(np.float64),
                            SAMPLE_RATE // g,
                            framerate // g,
                        ).astype(np.float32)

                    return audio
            except Exception as e:
                print(f"  [WAV decode error] {e}")
                return None

        def _send_error(self, code, message):
            self.send_response(code)
            self.send_header("Content-Type", "application/json")
            self.send_header("Access-Control-Allow-Origin", "*")
            self.end_headers()
            self.wfile.write(json.dumps({"error": message}).encode("utf-8"))

    return HTTPServer, RequestHandler


def main():
    parser = argparse.ArgumentParser(description="VoiceSwitch Test Server")
    parser.add_argument("--port", type=int, default=8080, help="Server port (default: 8080)")
    parser.add_argument("--model", type=str, default=MODEL_PATH, help="Path to model file")
    args = parser.parse_args()

    print("=" * 60)
    print("VoiceSwitch Test Server")
    print("=" * 60)

    predictor = VoiceSwitchPredictor(model_path=args.model)

    HTTPServer, Handler = create_app(predictor)
    server = HTTPServer(("0.0.0.0", args.port), Handler)

    print(f"\n  Server running at: http://localhost:{args.port}")
    print(f"  Open that URL in your browser to test the model.")
    print(f"  Press Ctrl+C to stop.\n")

    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\nShutting down...")
        server.shutdown()


if __name__ == "__main__":
    main()