vts_playback_driver.py

import os
import json
import asyncio
import numpy as np
import websockets

# Configuration
PORT = 8001
HOST = "ws://127.0.0.1"
NPY_PATH = "dataset\\landmarks\\XTRlxbOA8Pg_12fps_001153_to_end_norm_landmarks.npy"
TARGET_FPS = 12.0

# --- 28-point anime-face landmark map (ground-truthed from visualize_wireframe) ---
# 0  right face @ eye line   4  left face @ eye line
# 1  right face @ chin       3  left face @ chin       2  chin
# 5,6,7  right brow          8,9,10  left brow
# 11,12,13  left eye top (l,c,r)   14,15,16  left eye bottom (l,c,r)
# 17,18,19  right eye top (l,c,r)  20,21,22  right eye bottom (l,c,r)
# 23 nose tip
# 24 right mouth corner   26 left mouth corner   25 mouth top   27 mouth bottom
# (Y axis points down, image coords. "left"/"right" are the subject's, as drawn.)

# Calibration: resting-face baselines and spreads, measured from the 89k-frame
# dataset (percentiles). Each metric is centered on its median so a neutral
# expression maps to a neutral VTS value, then scaled by its spread.
MOUTH_OPEN_NEUTRAL = 0.02   # open_dist / mouth_width, mouth closed
MOUTH_OPEN_FULL    = 0.55   # ~p99 → MouthOpen=1.0
SMILE_NEUTRAL      = 0.085  # (center_y - corners_y) / mouth_width, neutral
SMILE_UP           = 0.185  # spread to MouthSmile=+1
SMILE_DOWN         = 0.20   # spread to MouthSmile=-1
YAW_NEUTRAL        = 0.267  # (Leye-nose - Reye-nose)/(sum), neutral
YAW_SPREAD         = 0.235  # → FaceAngleX=+30
PITCH_NEUTRAL      = 0.34   # (eye_y - nose_y)/(eye_y - chin_y), neutral
PITCH_SPREAD       = 0.171  # → FaceAngleY=+15 (up)
EYE_OPEN_FULL      = 0.15   # eyelid gap / face width → EyeOpen=1.0
EYE_CLOSED         = 0.01   # → EyeOpen=0.0
BROWL_NEUTRAL      = 0.64   # (eye_y - brow_y)/face_h, neutral
BROWR_NEUTRAL      = 0.59
BROW_UP            = 0.25   # spread to BrowY=+1
BROW_DOWN          = 0.38   # spread to BrowY=-1


def _clip(x, lo, hi):
    return max(lo, min(hi, x))


def calculate_params(lm):
    """Map one frame's 28 normalized landmarks to VTube Studio standard inputs.
    Returns None if tracking was lost (NaN) so the driver can drop the face."""
    if np.isnan(lm).any():
        return None

    # --- shared reference geometry ---
    eye_l = (lm[12] + lm[15]) / 2.0   # left eye center (top-center & bottom-center)
    eye_r = (lm[18] + lm[21]) / 2.0   # right eye center
    eye_y = (eye_l[1] + eye_r[1]) / 2.0
    face_w = np.linalg.norm(lm[0] - lm[4])
    mouth_w = np.linalg.norm(lm[24] - lm[26])

    # 1. MouthOpen — vertical gap of inner lips, normalized by mouth width.
    open_dist = np.linalg.norm(lm[25] - lm[27])
    mouth_open = _clip((open_dist / (mouth_w + 1e-6) - MOUTH_OPEN_NEUTRAL) /
                       (MOUTH_OPEN_FULL - MOUTH_OPEN_NEUTRAL), 0.0, 1.0)

    # 2. MouthSmile — corners raised vs mouth center, normalized by mouth width.
    corners_y = (lm[24, 1] + lm[26, 1]) / 2.0
    center_y = (lm[25, 1] + lm[27, 1]) / 2.0
    smile_r = (center_y - corners_y) / (mouth_w + 1e-6) - SMILE_NEUTRAL
    if smile_r >= 0:
        mouth_smile = _clip(smile_r / SMILE_UP, -1.0, 1.0)
    else:
        mouth_smile = _clip(smile_r / SMILE_DOWN, -1.0, 1.0)

    # 3. Head Roll (FaceAngleZ) — eye-line tilt.
    angle_z = np.degrees(np.arctan2(eye_r[1] - eye_l[1], eye_r[0] - eye_l[0]))

    # 4. Head Yaw (FaceAngleX) — nose offset between the two eyes.
    ld = np.linalg.norm(eye_l - lm[23])
    rd = np.linalg.norm(eye_r - lm[23])
    yaw = (ld - rd) / (ld + rd + 1e-6)
    angle_x = _clip((yaw - YAW_NEUTRAL) / YAW_SPREAD * 30.0, -30.0, 30.0)

    # 5. Head Pitch (FaceAngleY) — nose height between eye-line and chin.
    nose_y, chin_y = lm[23, 1], lm[2, 1]
    pitch_r = (eye_y - nose_y) / (eye_y - chin_y + 1e-6)
    angle_y = _clip((pitch_r - PITCH_NEUTRAL) / PITCH_SPREAD * 15.0, -15.0, 15.0)

    # 6. Eye openness — eyelid gap per eye, normalized by face width.
    l_top = (lm[11] + lm[12] + lm[13]) / 3.0
    l_bot = (lm[14] + lm[15] + lm[16]) / 3.0
    r_top = (lm[17] + lm[18] + lm[19]) / 3.0
    r_bot = (lm[20] + lm[21] + lm[22]) / 3.0
    l_open = _clip((np.linalg.norm(l_top - l_bot) / (face_w + 1e-6) - EYE_CLOSED) /
                   (EYE_OPEN_FULL - EYE_CLOSED), 0.0, 1.0)
    r_open = _clip((np.linalg.norm(r_top - r_bot) / (face_w + 1e-6) - EYE_CLOSED) /
                   (EYE_OPEN_FULL - EYE_CLOSED), 0.0, 1.0)

    # 7. Brows — brow height above eye line, normalized by face height.
    face_h = abs(chin_y - eye_y) + 1e-6
    browL = (eye_y - (lm[5, 1] + lm[6, 1] + lm[7, 1]) / 3.0) / face_h
    browR = (eye_y - (lm[8, 1] + lm[9, 1] + lm[10, 1]) / 3.0) / face_h
    browL_y = _clip((browL - BROWL_NEUTRAL) /
                    (BROW_UP if browL >= BROWL_NEUTRAL else BROW_DOWN), -1.0, 1.0)
    browR_y = _clip((browR - BROWR_NEUTRAL) /
                    (BROW_UP if browR >= BROWR_NEUTRAL else BROW_DOWN), -1.0, 1.0)

    return {
        "MouthOpen": float(mouth_open),
        "MouthSmile": float(mouth_smile),
        "FaceAngleX": float(angle_x),
        "FaceAngleY": float(angle_y),
        "FaceAngleZ": float(angle_z),
        "EyeOpenLeft": float(l_open),
        "EyeOpenRight": float(r_open),
        "BrowLeftY": float(browL_y),
        "BrowRightY": float(browR_y),
    }


SMOOTH_WINDOW = 3  # ponytail: centered moving average over consecutive valid frames;
# tames landmark jitter at 12 FPS without adding latency (offline playback, zero-phase).
# Bump to 5 if still choppy; each step trades responsiveness for smoothness.
def smooth_params(params):
    """Centered moving average over consecutive valid (non-None) frames.
    Gaps (None) act as boundaries — never averaged across — so a face-loss
    reset still snaps to neutral. Window is odd; edges shrink toward 1.
    Only cosmetic params are smoothed; angles/jaw pass through unchanged
    to preserve crisp head motion."""
    if SMOOTH_WINDOW <= 1:
        return params
    w = SMOOTH_WINDOW
    half = w // 2
    SMOOTH_KEYS = ("MouthOpen", "MouthSmile", "EyeOpenLeft", "EyeOpenRight",
                   "BrowLeftY", "BrowRightY")
    out = [None] * len(params)
    # Find runs of consecutive valid frames; smooth within each run.
    i = 0
    n = len(params)
    while i < n:
        if params[i] is None:
            i += 1
            continue
        j = i
        while j < n and params[j] is not None:
            j += 1
        # run = [i, j)
        for k in range(i, j):
            lo = max(i, k - half)
            hi = min(j, k + half + 1)
            avg = {key: 0.0 for key in SMOOTH_KEYS}
            for m in range(lo, hi):
                pt = params[m]
                for key in SMOOTH_KEYS:
                    avg[key] += pt[key]
            cnt = hi - lo
            merged = dict(params[k])
            for key in SMOOTH_KEYS:
                merged[key] = avg[key] / cnt
            out[k] = merged
        i = j
    return out


NEUTRAL_PARAMS = {
    "MouthOpen": 0.0, "MouthSmile": 0.0,
    "FaceAngleX": 0.0, "FaceAngleY": 0.0, "FaceAngleZ": 0.0,
    "EyeOpenLeft": 1.0, "EyeOpenRight": 1.0,
    "BrowLeftY": 0.0, "BrowRightY": 0.0,
}


NEUTRAL_PARAMS = {
    "MouthOpen": 0.0, "MouthSmile": 0.0,
    "FaceAngleX": 0.0, "FaceAngleY": 0.0, "FaceAngleZ": 0.0,
    "EyeOpenLeft": 1.0, "EyeOpenRight": 1.0,
    "BrowLeftY": 0.0, "BrowRightY": 0.0,
}


class CausalSmoother:
    """Trailing-window moving average for streaming. Same SMOOTH_KEYS as the
    offline centered smoother; only difference is causality (uses past frames
    only, no future) so it adds zero latency. A None param resets the run so a
    face-loss still snaps to neutral unsmoothed."""
    SMOOTH_KEYS = ("MouthOpen", "MouthSmile", "EyeOpenLeft", "EyeOpenRight",
                   "BrowLeftY", "BrowRightY")

    def __init__(self, window=SMOOTH_WINDOW):
        self.w = max(1, window)
        self.buf = []  # recent valid params in the current run

    def push(self, p):
        if p is None:
            self.buf = []
            return None
        self.buf.append(p)
        if len(self.buf) > self.w:
            self.buf = self.buf[-self.w:]
        merged = dict(p)
        for key in self.SMOOTH_KEYS:
            merged[key] = sum(b[key] for b in self.buf) / len(self.buf)
        return merged


async def vts_connect():
    """Open the websocket, do token handshake + auth + list params. Returns ws
    or None on failure."""
    print(f"Connecting to VTube Studio at {HOST}:{PORT}...")
    ws = await websockets.connect(f"{HOST}:{PORT}")
    auth_req = {
        "apiName": "VTubeStudioPublicAPI", "apiVersion": "1.0",
        "requestID": "auth_token_request", "messageType": "AuthenticationTokenRequest",
        "data": {"pluginName": "AnimeLandmarkDriver", "pluginDeveloper": "Dev"}}
    await ws.send(json.dumps(auth_req))
    resp = json.loads(await ws.recv())
    if "authenticationToken" not in resp.get("data", {}):
        print("ERROR: Authentication failed or timed out in VTube Studio UI.")
        await ws.close(); return None
    token = resp["data"]["authenticationToken"]
    login_req = {
        "apiName": "VTubeStudioPublicAPI", "apiVersion": "1.0",
        "requestID": "auth_request", "messageType": "AuthenticationRequest",
        "data": {"pluginName": "AnimeLandmarkDriver", "pluginDeveloper": "Dev",
                 "authenticationToken": token}}
    await ws.send(json.dumps(login_req))
    login_resp = json.loads(await ws.recv())
    if not login_resp["data"]["authenticated"]:
        print("Authentication rejected."); await ws.close(); return None
    print("Successfully authenticated.")
    # List params (informational; not required for injection).
    await ws.send(json.dumps({
        "apiName": "VTubeStudioPublicAPI", "apiVersion": "1.0",
        "requestID": "get_input_params", "messageType": "InputParameterListRequest"}))
    input_resp = json.loads(await ws.recv())
    if "data" in input_resp:
        dp = input_resp["data"].get("defaultParameters", [])
        cp = input_resp["data"].get("customParameters", [])
        dl = [p.get("name", p.get("id", "?")) for p in dp]
        cl = [p.get("name", p.get("id", "?")) for p in cp]
        print(f"Default Parameters: {', '.join(dl) if len(dl) <= 60 else ', '.join(dl[:60]) + '...'}")
        print(f"Custom Parameters: {', '.join(cl) if cl else 'None'}")
    return ws


async def vts_send_frame(ws, p):
    """Inject one param frame (p: dict or None=face lost)."""
    if p is None:
        data = {"faceFound": False, "mode": "set",
                "parameterValues": [{"id": k, "value": v} for k, v in NEUTRAL_PARAMS.items()]}
    else:
        data = {"faceFound": True, "mode": "set",
                "parameterValues": [{"id": k, "value": v} for k, v in p.items()]}
    await ws.send(json.dumps({
        "apiName": "VTubeStudioPublicAPI", "apiVersion": "1.0",
        "requestID": "inject_motion", "messageType": "InjectParameterDataRequest",
        "data": data}))
    await ws.recv()  # Clear WebSocket buffer


async def vts_release(ws):
    """Send one neutral frame so VTS stops holding the last pose."""
    await vts_send_frame(ws, None)


# ponytail: the NAT model compresses parameter variance by 3-5x vs real
# landmarks (MSE regression-to-mean on a one-to-many task). Gains restore
# visible dynamic range. Keep at 2x until the mouth-weighted loss (in
# 0.3_train_nat.py) trains a model with natural variance; then drop to 1x.
MODEL_GAIN = {
    "MouthOpen":    2.0,
    "MouthSmile":   2.0,
    "EyeOpenLeft":  2.0,
    "EyeOpenRight": 2.0,
    "BrowLeftY":    2.0,
    "BrowRightY":   2.0,
}

def _apply_gain(p):
    """Amplify compressed params back toward real-landmark variance."""
    if p is None:
        return None
    for key, gain in MODEL_GAIN.items():
        if key in p:
            p[key] = max(-1.0, min(1.0, p[key] * gain))
    return p


async def vts_driver_stream(get_frame, ws=None):
    """Streaming driver. `get_frame` is an awaitable callable returning the next
    landmark frame [28,2] or None (face lost); raise StopAsyncError / return a
    sentinel to end the stream. Paces output at TARGET_FPS, applies causal
    smoothing. This is the real-time path: compute for chunk N overlaps with
    playback of chunk N-1.

    If `ws` is given (pre-connected), skip the handshake and use it directly.
    This lets the caller connect before a user prompt so audio and animation
    start truly simultaneously."""
    if ws is None:
        ws = await vts_connect()
        if ws is None:
            return
    smoother = CausalSmoother(SMOOTH_WINDOW)
    delay = 1.0 / TARGET_FPS
    loop = asyncio.get_event_loop()
    print("Starting motion injection (streaming)...")
    try:
        # Steady-schedule pacing: target one send every `delay`, NOT `delay -
        # get_frame_wait`. The old code counted the queue-wait toward `elapsed`,
        # so after a producer stall it skipped the sleep and dumped the next
        # frame immediately -- two frames back-to-back (bursts of ~0.2ms gaps),
        # which VTS would partially drop ('API sent at the same time').
        # Resync (no catchup burst) if we ever fall behind.
        next_send = loop.time()
        while True:
            lm = await get_frame()
            if lm is _STREAM_END:
                break
            p = calculate_params(lm) if (lm is not None and not np.isnan(lm).any()) else None
            p = _apply_gain(p)
            p = smoother.push(p)
            now = loop.time()
            if now < next_send:
                await asyncio.sleep(next_send - now)
            await vts_send_frame(ws, p)
            next_send += delay
            if next_send < loop.time():  # fell behind: resync, don't pile up
                next_send = loop.time()
    finally:
        await vts_release(ws)
        await ws.close()
        print("Playback finished.")


_STREAM_END = object()  # sentinel: producer signals end-of-stream


async def vts_driver(landmarks=None):
    # `landmarks`: either pass a precomputed [V, 28, 2] array OR leave None to
    # load the dataset NPY_PATH. Offline path: precompute + centered smoothing.
    if landmarks is None:
        landmarks = np.load(NPY_PATH)
    raw_params = [calculate_params(lm) for lm in landmarks]
    raw_params = smooth_params(raw_params)
    ws = await vts_connect()
    if ws is None:
        return
    delay = 1.0 / TARGET_FPS
    loop = asyncio.get_event_loop()
    print("Starting motion injection...")
    try:
        for p in raw_params:
            t0 = loop.time()
            await vts_send_frame(ws, p)
            elapsed = loop.time() - t0
            if elapsed < delay:
                await asyncio.sleep(delay - elapsed)
    finally:
        await vts_release(ws)
        await ws.close()
        print("Playback finished.")


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