Mr-Saab29/poseflow_backend

1// src/pose.ts2import * as tf from '@tensorflow/tfjs';3import * as posedetection from '@tensorflow-models/pose-detection';4import { createCanvas, loadImage, Image } from 'canvas';5import fetch from 'node-fetch'; // <-- ensure this is in package.json6 7import { POINTS, CLASS_NO } from './constants.js';8 9/** ===== Backend init (CPU) =====10 * CPU backend is the most portable on Spaces. It's slower than WASM/WebGL,11 * but with LIGHTNING + downscaling it stays well within the request timeout.12 */13let backendReady: Promise<void> | null = null;14export async function ensureWasmBackend() {15  if (!backendReady) {16    backendReady = (async () => {17      await tf.setBackend('cpu');18      await tf.ready();19    })();20  }21  return backendReady;22}23 24// One-shot initializer you can call at startup/warmup25export async function initTf() {26  await ensureWasmBackend();27  await Promise.all([loadModel(), getDetector()]);28}29 30/** ===== Classifier model (unchanged) ===== */31const MODEL_URL =32  process.env.POSE_CLASSIFIER_URL ??33  'https://models.s3.jp-tok.cloud-object-storage.appdomain.cloud/model.json';34 35let modelPromise: Promise<tf.LayersModel> | null = null;36export async function loadModel(): Promise<tf.LayersModel> {37  if (!modelPromise) {38    modelPromise = tf.loadLayersModel(MODEL_URL);39  }40  return modelPromise;41}42 43/** ===== MoveNet (use LIGHTNING for speed) ===== */44let detectorPromise: Promise<posedetection.PoseDetector> | null = null;45async function getDetector(): Promise<posedetection.PoseDetector> {46  if (!detectorPromise) {47    await ensureWasmBackend();48    detectorPromise = posedetection.createDetector(49      posedetection.SupportedModels.MoveNet,50      {51        modelType: posedetection.movenet.modelType.SINGLEPOSE_LIGHTNING, // faster than THUNDER52      }53    );54  }55  return detectorPromise;56}57 58/** ===== Helpers ===== */59 60// Fetch an image with a timeout, then load via canvas61async function fetchImageWithTimeout(url: string, ms = 8000): Promise<Image> {62  const ctrl = new AbortController();63  const id = setTimeout(() => ctrl.abort(), ms);64  try {65    const r = await fetch(url, { signal: ctrl.signal });66    if (!r.ok) throw new Error(`HTTP ${r.status}`);67    const buf = Buffer.from(await r.arrayBuffer());68    return await loadImage(buf);69  } finally {70    clearTimeout(id);71  }72}73 74// Downscale while keeping aspect ratio (important for speed)75function toResizedCanvas(img: Image, maxSide = 512) {76  let { width: w, height: h } = img;77  const scale = Math.min(1, maxSide / Math.max(w, h));78  const W = Math.max(1, Math.round(w * scale));79  const H = Math.max(1, Math.round(h * scale));80  const canvas = createCanvas(W, H);81  const ctx = canvas.getContext('2d');82  ctx.drawImage(img, 0, 0, W, H);83  return canvas;84}85 86// Image/B64 -> Canvas (downscaled)87async function loadImageToCanvas(input: { imageUrl?: string; imageB64?: string }) {88  let img: Image;89  if (input.imageUrl) {90    img = await fetchImageWithTimeout(input.imageUrl, 8000);91  } else if (input.imageB64) {92    const buf = Buffer.from(input.imageB64, 'base64');93    img = await loadImage(buf);94  } else {95    throw new Error('Provide imageUrl or imageB64');96  }97  return toResizedCanvas(img, 512);98}99 100// Promise timeout helper101function withTimeout<T>(p: Promise<T>, ms = 8000, msg = 'Timed out') {102  return Promise.race<T>([103    p,104    new Promise<T>((_, rej) => setTimeout(() => rej(new Error(msg)), ms)),105  ]);106}107 108/** ===== Image -> 17×2 normalized landmarks ===== */109export async function detectLandmarksFromImage(input: { imageUrl?: string; imageB64?: string }): Promise<number[][]> {110  const canvas = await loadImageToCanvas(input);111  const detector = await getDetector();112 113  // Guard estimatePoses with an 8s timeout114  const poses = await withTimeout(115    detector.estimatePoses(canvas as any, { flipHorizontal: false }),116    8000,117    'pose detection timed out'118  );119 120  if (!poses.length || !poses[0].keypoints || poses[0].keypoints.length < 17) {121    throw new Error('No pose detected');122  }123 124  const kp = poses[0].keypoints.slice(0, 17);125  const w = canvas.width, h = canvas.height;126  const landmarks = kp.map((k) => [k.x / w, k.y / h]);127  return landmarks;128}129 130/** ===== Preprocess + Classify (same as before) ===== */131function get_center_point(landmarks: tf.Tensor2D, leftIdx: number, rightIdx: number) {132  const left = tf.gather(landmarks, leftIdx, 0);133  const right = tf.gather(landmarks, rightIdx, 0);134  return tf.add(tf.mul(left, 0.5), tf.mul(right, 0.5));135}136 137function get_pose_size(landmarks: tf.Tensor2D, torsoSizeMultiplier = 2.5) {138  const hipsCenter = get_center_point(landmarks, POINTS.LEFT_HIP, POINTS.RIGHT_HIP);139  const shouldersCenter = get_center_point(landmarks, POINTS.LEFT_SHOULDER, POINTS.RIGHT_SHOULDER);140  const torsoSize = tf.norm(tf.sub(shouldersCenter, hipsCenter));141 142  let poseCenter = get_center_point(landmarks, POINTS.LEFT_HIP, POINTS.RIGHT_HIP);143  poseCenter = tf.expandDims(poseCenter, 0);144  poseCenter = tf.tile(poseCenter, [17, 1]);145 146  const d = tf.sub(landmarks, poseCenter);147  const maxDist = tf.max(tf.norm(d, 'euclidean', 1));148  return tf.maximum(tf.mul(torsoSize, torsoSizeMultiplier), maxDist);149}150 151function normalize_pose_landmarks(landmarks: tf.Tensor2D) {152  let poseCenter = get_center_point(landmarks, POINTS.LEFT_HIP, POINTS.RIGHT_HIP);153  poseCenter = tf.expandDims(poseCenter, 0);154  poseCenter = tf.tile(poseCenter, [17, 1]);155  const shifted = tf.sub(landmarks, poseCenter);156  const poseSize = get_pose_size(landmarks);157  return tf.div(shifted, poseSize) as tf.Tensor2D;158}159 160function landmarks_to_embedding(landmarks: tf.Tensor2D) {161  const norm = normalize_pose_landmarks(landmarks);162  return tf.reshape(norm, [1, 34]);163}164 165export async function classifyLandmarks(landmarksArray: number[][], threshold = 0.8) {166  await ensureWasmBackend(); // make sure backend is ready167 168  if (!Array.isArray(landmarksArray) || landmarksArray.length !== 17) {169    throw new Error('landmarks must be 17x2 array');170  }171  const model = await loadModel();172 173  const landmarks = tf.tensor2d(landmarksArray, [17, 2]);174  const input = landmarks_to_embedding(landmarks);175  const logits = model.predict(input) as tf.Tensor;176  const scores = (await logits.array()) as number[][];177  input.dispose(); logits.dispose(); landmarks.dispose();178 179  const probs = scores[0];180  const entries = Object.entries(CLASS_NO).map(([name, idx]) => [name, probs[idx] ?? 0] as const);181  entries.sort((a, b) => b[1] - a[1]);182 183  const [bestName, bestScore] = entries[0];184  return {185    top: { name: bestName, score: bestScore },186    scores: Object.fromEntries(entries),187    passed: bestScore >= threshold,188    threshold,189  };190}191 192// For compatibility with older callers193export async function predictPose(input: { imageUrl?: string; imageB64?: string }) {194  const landmarks = await detectLandmarksFromImage(input);195  return classifyLandmarks(landmarks, 0.8);196}