Project_Velocity/app/src/hooks/useMediapipeFaceLandmarker.ts

import { useState, useEffect, useRef, useCallback } from 'react';
import {
  FaceLandmarker,
  FilesetResolver,
} from '@mediapipe/tasks-vision';

export interface BlendShapeCategory {
  categoryName: string;
  score: number;
  displayName: string;
  index: number;
}

export interface FaceLandmarkerResult {
  faceBlendshapes: Array<{ categories: BlendShapeCategory[] }>;
  faceLandmarks: Array<Array<{ x: number; y: number; z: number }>>;
}

const MODEL_URL =
  import.meta.env.VITE_MEDIAPIPE_MODEL_URL ??
  '/mediapipe/assets/face_landmarker.task';

const WASM_ROOT =
  import.meta.env.VITE_MEDIAPIPE_WASM_ROOT ??
  '/mediapipe/wasm';

interface UseFaceLandmarkerReturn {
  isLoading: boolean;
  isReady: boolean;
  error: string | null;
  detectFrame: (
    videoElement: HTMLVideoElement,
    timestampMs: number,
  ) => FaceLandmarkerResult | null;
}

export function useMediapipeFaceLandmarker(): UseFaceLandmarkerReturn {
  const landmarkerRef = useRef<FaceLandmarker | null>(null);
  const [isLoading, setIsLoading] = useState(true);
  const [isReady, setIsReady] = useState(false);
  const [error, setError] = useState<string | null>(null);

  useEffect(() => {
    let cancelled = false;

    async function init() {
      try {
        const filesetResolver = await FilesetResolver.forVisionTasks(WASM_ROOT);
        const landmarker = await FaceLandmarker.createFromOptions(filesetResolver, {
          baseOptions: {
            modelAssetPath: MODEL_URL,
            delegate: 'GPU',
          },
          outputFaceBlendshapes: true,
          runningMode: 'VIDEO',
          numFaces: 3,
          minFaceDetectionConfidence: 0.65,
          minFacePresenceConfidence: 0.6,
          minTrackingConfidence: 0.6,
        });

        if (cancelled) {
          landmarker.close();
          return;
        }

        landmarkerRef.current = landmarker;
        setIsLoading(false);
        setIsReady(true);
      } catch (err) {
        if (cancelled) return;
        console.error('[MediaPipe] Initialization failed:', err);
        setError(err instanceof Error ? err.message : 'MediaPipe failed to initialize.');
        setIsLoading(false);
      }
    }

    void init();

    return () => {
      cancelled = true;
      landmarkerRef.current?.close();
      landmarkerRef.current = null;
    };
  }, []);

  const detectFrame = useCallback(
    (videoElement: HTMLVideoElement, timestampMs: number): FaceLandmarkerResult | null => {
      if (!landmarkerRef.current || !isReady) return null;
      try {
        return landmarkerRef.current.detectForVideo(videoElement, timestampMs) as FaceLandmarkerResult;
      } catch {
        return null;
      }
    },
    [isReady],
  );

  return { isLoading, isReady, error, detectFrame };
}