214 lines
8.0 KiB
Python
214 lines
8.0 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import math
|
|
import time
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
|
|
from config import settings
|
|
|
|
|
|
class AvatarService:
|
|
def __init__(self) -> None:
|
|
self.driver_name = "arkit-blendshape-v1"
|
|
self._lock = asyncio.Lock()
|
|
self._sequence = 0
|
|
self._last_error: str | None = None
|
|
self._last_frame_count = 0
|
|
self._last_chunk_duration_ms = 0
|
|
|
|
async def build_idle_payload(self, reason: str = "idle") -> dict[str, Any]:
|
|
frame = self._build_frame(seq=await self._next_sequence(), time_ms=0, controls=self._neutral_controls())
|
|
self._last_frame_count = 1
|
|
self._last_chunk_duration_ms = 0
|
|
self._last_error = None
|
|
return {
|
|
"type": "animation_state",
|
|
"driver": self.driver_name,
|
|
"mode": settings.avatar_driver_mode,
|
|
"schema": settings.avatar_blendshape_schema,
|
|
"fps": settings.avatar_fps,
|
|
"reason": reason,
|
|
"frame_count": 1,
|
|
"frames": [frame],
|
|
"ts_ms": int(time.time() * 1000),
|
|
}
|
|
|
|
async def build_reset_payload(self, reason: str = "reset") -> dict[str, Any]:
|
|
payload = await self.build_idle_payload(reason=reason)
|
|
payload["type"] = "animation_reset"
|
|
return payload
|
|
|
|
async def build_controls_from_audio(
|
|
self,
|
|
audio: np.ndarray,
|
|
sr: int,
|
|
*,
|
|
text: str = "",
|
|
chunk_index: int = 0,
|
|
total_chunks: int = 1,
|
|
) -> dict[str, Any]:
|
|
if sr <= 0:
|
|
self._last_error = f"invalid sample rate: {sr}"
|
|
return await self.build_idle_payload(reason="invalid-sample-rate")
|
|
|
|
mono = np.asarray(audio, dtype=np.float32).reshape(-1)
|
|
if mono.size <= 1:
|
|
return await self.build_idle_payload(reason="empty-audio")
|
|
|
|
mono = np.clip(mono, -1.0, 1.0)
|
|
frame_count = max(1, int(math.ceil((mono.shape[0] / float(sr)) * settings.avatar_fps)))
|
|
seqs = await self._reserve_sequences(frame_count)
|
|
frames = self._audio_to_frames(mono, sr, seqs)
|
|
|
|
self._last_frame_count = len(frames)
|
|
self._last_chunk_duration_ms = int(round(1000.0 * mono.shape[0] / float(sr)))
|
|
self._last_error = None
|
|
return {
|
|
"type": "animation_chunk",
|
|
"driver": self.driver_name,
|
|
"mode": settings.avatar_driver_mode,
|
|
"schema": settings.avatar_blendshape_schema,
|
|
"chunk_index": chunk_index,
|
|
"total_chunks": total_chunks,
|
|
"sample_rate": sr,
|
|
"fps": settings.avatar_fps,
|
|
"text": text,
|
|
"frame_count": len(frames),
|
|
"duration_ms": self._last_chunk_duration_ms,
|
|
"frames": frames,
|
|
"ts_ms": int(time.time() * 1000),
|
|
}
|
|
|
|
@property
|
|
def schema(self) -> dict[str, Any]:
|
|
return {
|
|
"driver": self.driver_name,
|
|
"mode": settings.avatar_driver_mode,
|
|
"protocol": settings.avatar_control_protocol,
|
|
"schema": settings.avatar_blendshape_schema,
|
|
"controls": [
|
|
"jawOpen",
|
|
"mouthClose",
|
|
"mouthFunnel",
|
|
"mouthPucker",
|
|
"viseme_aa",
|
|
"viseme_ee",
|
|
"viseme_oh",
|
|
"headYaw",
|
|
"headPitch",
|
|
"headRoll",
|
|
],
|
|
}
|
|
|
|
def _audio_to_frames(self, audio: np.ndarray, sr: int, sequences: range) -> list[dict[str, Any]]:
|
|
samples_per_frame = max(1, int(round(sr / max(1, settings.avatar_fps))))
|
|
frames: list[dict[str, Any]] = []
|
|
jaw_prev = 0.0
|
|
for index, seq in enumerate(sequences):
|
|
start = index * samples_per_frame
|
|
end = min(audio.shape[0], start + samples_per_frame)
|
|
window = audio[start:end]
|
|
if window.size == 0:
|
|
window = np.zeros(1, dtype=np.float32)
|
|
energy, brightness = self._extract_features(window, sr)
|
|
jaw_open = 0.70 * jaw_prev + 0.30 * energy
|
|
jaw_prev = jaw_open
|
|
phase = index / max(1, len(sequences) - 1)
|
|
frames.append(
|
|
self._build_frame(
|
|
seq=seq,
|
|
time_ms=int(round(index * 1000.0 / settings.avatar_fps)),
|
|
controls=self._controls_from_features(jaw_open, brightness, phase),
|
|
)
|
|
)
|
|
return frames
|
|
|
|
def _extract_features(self, window: np.ndarray, sr: int) -> tuple[float, float]:
|
|
rms = float(np.sqrt(np.mean(window**2))) if window.size else 0.0
|
|
energy = float(np.clip((rms - 0.01) * 10.0, 0.0, 1.0))
|
|
if window.size < 16:
|
|
return energy, 0.5
|
|
spectrum = np.abs(np.fft.rfft(window * np.hanning(window.size)))
|
|
if spectrum.size <= 1:
|
|
return energy, 0.5
|
|
freqs = np.fft.rfftfreq(window.size, d=1.0 / float(sr))
|
|
weights = spectrum[1:]
|
|
freq_values = freqs[1:]
|
|
denom = float(np.sum(weights)) + 1e-6
|
|
centroid = float(np.sum(freq_values * weights) / denom)
|
|
brightness = float(np.clip((centroid - 500.0) / 2500.0, 0.0, 1.0))
|
|
return energy, brightness
|
|
|
|
def _controls_from_features(self, energy: float, brightness: float, phase: float) -> dict[str, float]:
|
|
jaw_open = float(np.clip(energy, 0.0, 1.0))
|
|
mouth_close = float(np.clip(1.0 - jaw_open * 0.82, 0.0, 1.0))
|
|
mouth_funnel = float(np.clip(0.16 + jaw_open * (0.40 - brightness * 0.12), 0.0, 1.0))
|
|
mouth_pucker = float(np.clip(0.06 + jaw_open * (0.20 + (1.0 - brightness) * 0.30), 0.0, 1.0))
|
|
viseme_aa = jaw_open
|
|
viseme_ee = float(np.clip(jaw_open * (0.25 + brightness * 0.95), 0.0, 1.0))
|
|
viseme_oh = float(np.clip(jaw_open * (0.30 + (1.0 - brightness) * 0.90), 0.0, 1.0))
|
|
head_yaw = float(np.clip(math.sin(phase * math.pi * 1.7) * 0.02 * max(jaw_open, 0.18), -0.25, 0.25))
|
|
head_pitch = float(np.clip(math.cos(phase * math.pi * 2.0) * 0.03 * max(jaw_open, 0.15), -0.25, 0.25))
|
|
head_roll = float(np.clip(math.sin(phase * math.pi) * 0.015, -0.15, 0.15))
|
|
return {
|
|
"jawOpen": round(jaw_open, 4),
|
|
"mouthClose": round(mouth_close, 4),
|
|
"mouthFunnel": round(mouth_funnel, 4),
|
|
"mouthPucker": round(mouth_pucker, 4),
|
|
"viseme_aa": round(viseme_aa, 4),
|
|
"viseme_ee": round(viseme_ee, 4),
|
|
"viseme_oh": round(viseme_oh, 4),
|
|
"headYaw": round(head_yaw, 4),
|
|
"headPitch": round(head_pitch, 4),
|
|
"headRoll": round(head_roll, 4),
|
|
}
|
|
|
|
def _build_frame(self, seq: int, time_ms: int, controls: dict[str, float]) -> dict[str, Any]:
|
|
return {
|
|
"seq": seq,
|
|
"time_ms": time_ms,
|
|
"speaking": controls["jawOpen"] > 0.04,
|
|
"controls": controls,
|
|
}
|
|
|
|
def _neutral_controls(self) -> dict[str, float]:
|
|
return {
|
|
"jawOpen": 0.0,
|
|
"mouthClose": 1.0,
|
|
"mouthFunnel": 0.0,
|
|
"mouthPucker": 0.0,
|
|
"viseme_aa": 0.0,
|
|
"viseme_ee": 0.0,
|
|
"viseme_oh": 0.0,
|
|
"headYaw": 0.0,
|
|
"headPitch": 0.0,
|
|
"headRoll": 0.0,
|
|
}
|
|
|
|
async def _next_sequence(self) -> int:
|
|
async with self._lock:
|
|
seq = self._sequence
|
|
self._sequence += 1
|
|
return seq
|
|
|
|
async def _reserve_sequences(self, count: int) -> range:
|
|
async with self._lock:
|
|
start = self._sequence
|
|
self._sequence += count
|
|
return range(start, start + count)
|
|
|
|
@property
|
|
def health(self) -> dict[str, Any]:
|
|
return {
|
|
"driver": self.driver_name,
|
|
"mode": settings.avatar_driver_mode,
|
|
"protocol": settings.avatar_control_protocol,
|
|
"schema": settings.avatar_blendshape_schema,
|
|
"last_frame_count": self._last_frame_count,
|
|
"last_chunk_duration_ms": self._last_chunk_duration_ms,
|
|
"last_error": self._last_error,
|
|
}
|