product/scripts/smoke_test.py

from __future__ import annotations

import argparse
import asyncio
import math
import os
import sys
import wave

import cv2
import numpy as np

PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if PROJECT_ROOT not in sys.path:
    sys.path.insert(0, PROJECT_ROOT)

from core.pipeline import ChatPipeline
from core.state_machine import Arbitrator
from services.asr import ASRService
from services.avatar import AvatarService
from services.llm import LLMService
from services.tts import TTSService


def _write_wav(path: str, audio: np.ndarray, sr: int) -> None:
    pcm = np.clip(audio * 32767.0, -32768, 32767).astype(np.int16)
    with wave.open(path, "wb") as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(sr)
        wf.writeframes(pcm.tobytes())


async def _write_avatar_video(
    avatar: AvatarService, audio: np.ndarray, sr: int, out_path: str, fps: int
) -> tuple[int, float]:
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    width = 720
    height = 720
    writer = cv2.VideoWriter(out_path, fourcc, float(fps), (width, height))
    if not writer.isOpened():
        raise RuntimeError(f"failed to open video writer: {out_path}")

    payload = await avatar.build_controls_from_audio(audio, sr, text="smoke-test")
    frames = payload.get("frames", [])
    duration_ms = int(payload.get("duration_ms", 0))

    for frame_payload in frames:
        controls = frame_payload.get("controls", {})
        frame = _render_debug_frame(width, height, controls, speaking=bool(frame_payload.get("speaking", False)))
        writer.write(frame)

    # add 0.4s idle tail for easy visual check
    tail_frames = max(1, int(0.4 * fps))
    for _ in range(tail_frames):
        frame = _render_debug_frame(width, height, {}, speaking=False)
        writer.write(frame)

    writer.release()
    return len(frames) + tail_frames, (duration_ms / 1000.0) + 0.4


def _render_debug_frame(width: int, height: int, controls: dict, speaking: bool) -> np.ndarray:
    jaw_open = float(np.clip(controls.get("jawOpen", 0.0), 0.0, 1.0))
    mouth_pucker = float(np.clip(controls.get("mouthPucker", 0.0), 0.0, 1.0))
    viseme_aa = float(np.clip(controls.get("viseme_aa", jaw_open), 0.0, 1.0))
    viseme_ee = float(np.clip(controls.get("viseme_ee", 0.0), 0.0, 1.0))
    viseme_oh = float(np.clip(controls.get("viseme_oh", 0.0), 0.0, 1.0))
    head_yaw = float(np.clip(controls.get("headYaw", 0.0), -1.0, 1.0))
    head_pitch = float(np.clip(controls.get("headPitch", 0.0), -1.0, 1.0))
    head_roll = float(np.clip(controls.get("headRoll", 0.0), -1.0, 1.0))

    frame = np.zeros((height, width, 3), dtype=np.uint8)
    frame[:] = (16, 12, 10)

    center_x = width // 2 + int(head_yaw * 48)
    center_y = height // 2 + int(head_pitch * 42) - 24
    head_radius = 190

    cv2.circle(frame, (center_x, center_y), head_radius + 26, (42, 28, 18), -1, lineType=cv2.LINE_AA)
    cv2.circle(frame, (center_x, center_y), head_radius, (198, 214, 238), -1, lineType=cv2.LINE_AA)

    eye_y = center_y - 44
    eye_offset = 64
    cv2.circle(frame, (center_x - eye_offset, eye_y), 12, (26, 31, 43), -1, lineType=cv2.LINE_AA)
    cv2.circle(frame, (center_x + eye_offset, eye_y), 12, (26, 31, 43), -1, lineType=cv2.LINE_AA)

    brow_dx = int(head_roll * 28)
    cv2.line(frame, (center_x - 100 - brow_dx, eye_y - 34), (center_x - 42 - brow_dx, eye_y - 40), (48, 58, 74), 6, lineType=cv2.LINE_AA)
    cv2.line(frame, (center_x + 42 - brow_dx, eye_y - 40), (center_x + 100 - brow_dx, eye_y - 34), (48, 58, 74), 6, lineType=cv2.LINE_AA)

    mouth_center = (center_x, center_y + 86)
    mouth_width = int(64 + viseme_oh * 36 + mouth_pucker * 18)
    mouth_height = int(12 + jaw_open * 98 + viseme_aa * 18)
    lip_color = (96, 54, 130) if speaking else (88, 74, 108)
    inner_color = (36, 22, 54)
    cv2.ellipse(frame, mouth_center, (mouth_width, max(6, mouth_height)), 0, 0, 360, lip_color, -1, lineType=cv2.LINE_AA)
    cv2.ellipse(frame, mouth_center, (max(12, mouth_width - 16), max(4, mouth_height - 8)), 0, 0, 360, inner_color, -1, lineType=cv2.LINE_AA)

    ee_width = int(max(20, 86 - viseme_ee * 46))
    cv2.line(frame, (mouth_center[0] - ee_width, mouth_center[1]), (mouth_center[0] + ee_width, mouth_center[1]), (220, 224, 232), 3, lineType=cv2.LINE_AA)

    metrics = [
        f"jawOpen={jaw_open:.2f}",
        f"viseme_aa={viseme_aa:.2f}",
        f"viseme_ee={viseme_ee:.2f}",
        f"viseme_oh={viseme_oh:.2f}",
        f"headYaw={head_yaw:.2f}",
    ]
    for index, text in enumerate(metrics):
        cv2.putText(frame, text, (28, 42 + index * 28), cv2.FONT_HERSHEY_SIMPLEX, 0.72, (210, 220, 235), 2, lineType=cv2.LINE_AA)

    return frame


async def main() -> None:
    parser = argparse.ArgumentParser(description="Offline full-pipeline smoke test")
    parser.add_argument("--text", default="你好，请做一个十秒内的简短自我介绍。")
    parser.add_argument("--out-dir", default="outputs")
    parser.add_argument("--fps", type=int, default=25)
    args = parser.parse_args()

    os.makedirs(args.out_dir, exist_ok=True)
    wav_path = os.path.join(args.out_dir, "smoke_reply.wav")
    mp4_path = os.path.join(args.out_dir, "smoke_avatar.mp4")

    arbitrator = Arbitrator()
    pipeline = ChatPipeline(
        arbitrator=arbitrator,
        asr=ASRService(),
        llm=LLMService(),
        tts=TTSService(),
    )
    avatar = AvatarService()

    result, audio, sr = await pipeline.process_text_turn(args.text)
    _write_wav(wav_path, audio, sr)
    frame_count, video_s = await _write_avatar_video(avatar, audio, sr, mp4_path, args.fps)

    print("=== smoke test done ===")
    print(f"user_text: {result.get('user_text', '')}")
    print(f"reply_text: {result.get('reply_text', '')}")
    print(f"llm_source: {result.get('llm_source', 'unknown')}")
    print(f"audio_samples: {result.get('audio_samples', 0)}, sample_rate: {sr}")
    print(f"video_frames: {frame_count}, video_seconds: {video_s:.2f}")
    print(f"wav: {wav_path}")
    print(f"mp4: {mp4_path}")


if __name__ == "__main__":
    asyncio.run(main())