rtmlib-ts 0.0.2

package/src/solution/poseTracker.ts

@@ -0,0 +1,172 @@
+/**
+ * PoseTracker - tracks poses across frames with cached detections
+ * Reduces detection frequency for better performance
+ */
+
+import { Wholebody } from './wholebody';
+import { BBox } from '../types/index';
+
+interface TrackedBox {
+  bbox: BBox;
+  age: number;
+  lastSeen: number;
+  id: number;
+}
+
+export class PoseTracker {
+  private wholebody: Wholebody;
+  private detFrequency: number;
+  private cachedBoxes: TrackedBox[] = [];
+  private frameCount: number = 0;
+  private nextId: number = 0;
+
+  constructor(
+    WholebodyClass: typeof Wholebody,
+    detFrequency: number = 7,
+    toOpenpose: boolean = false,
+    mode: 'performance' | 'lightweight' | 'balanced' = 'balanced',
+    backend: 'onnxruntime' = 'onnxruntime',
+    device: string = 'cpu'
+  ) {
+    this.detFrequency = detFrequency;
+    this.wholebody = new WholebodyClass(
+      null,
+      [640, 640],
+      null,
+      [288, 384],
+      mode,
+      toOpenpose,
+      backend,
+      device
+    );
+  }
+
+  async init(): Promise<void> {
+    await this.wholebody.init();
+  }
+
+  async call(
+    image: Uint8Array,
+    imgWidth: number,
+    imgHeight: number
+  ): Promise<{ keypoints: number[][]; scores: number[] }> {
+    this.frameCount++;
+
+    // Run detection periodically
+    if (this.frameCount % this.detFrequency === 0 || this.cachedBoxes.length === 0) {
+      const result = await this.wholebody.call(image, imgWidth, imgHeight);
+      this.updateCachedBoxes(result.keypoints, result.scores, imgWidth, imgHeight);
+    }
+
+    // Use cached boxes for pose estimation
+    const bboxes = this.cachedBoxes.map((tb) => tb.bbox);
+    const result = await this.wholebody.call(image, imgWidth, imgHeight, bboxes);
+
+    // Clean up old boxes
+    this.cleanupCachedBoxes();
+
+    return result;
+  }
+
+  private updateCachedBoxes(
+    keypoints: number[][],
+    scores: number[],
+    imgWidth: number,
+    imgHeight: number
+  ): void {
+    // Simple tracking: create new boxes from keypoints
+    const newBoxes: TrackedBox[] = [];
+
+    for (let i = 0; i < keypoints.length; i += 17) {
+      const instanceKeypoints = keypoints.slice(i, Math.min(i + 17, keypoints.length));
+      const instanceScores = scores.slice(i, Math.min(i + 17, scores.length));
+
+      // Calculate bounding box from keypoints
+      let minX = imgWidth;
+      let minY = imgHeight;
+      let maxX = 0;
+      let maxY = 0;
+
+      for (let j = 0; j < instanceKeypoints.length; j++) {
+        if (instanceScores[j] > 0.3) {
+          const [x, y] = instanceKeypoints[j];
+          minX = Math.min(minX, x);
+          minY = Math.min(minY, y);
+          maxX = Math.max(maxX, x);
+          maxY = Math.max(maxY, y);
+        }
+      }
+
+      if (maxX > minX && maxY > minY) {
+        // Add padding
+        const padding = 0.25;
+        const width = maxX - minX;
+        const height = maxY - minY;
+        const paddedWidth = width * (1 + padding);
+        const paddedHeight = height * (1 + padding);
+
+        const centerX = (minX + maxX) / 2;
+        const centerY = (minY + maxY) / 2;
+
+        const x1 = Math.max(0, centerX - paddedWidth / 2);
+        const y1 = Math.max(0, centerY - paddedHeight / 2);
+        const x2 = Math.min(imgWidth, centerX + paddedWidth / 2);
+        const y2 = Math.min(imgHeight, centerY + paddedHeight / 2);
+
+        // Try to match with existing boxes
+        let matched = false;
+        for (const cachedBox of this.cachedBoxes) {
+          const iou = this.calculateIoU(
+            [x1, y1, x2, y2],
+            [cachedBox.bbox.x1, cachedBox.bbox.y1, cachedBox.bbox.x2, cachedBox.bbox.y2]
+          );
+
+          if (iou > 0.3) {
+            cachedBox.bbox = { x1, y1, x2, y2 };
+            cachedBox.age = 0;
+            cachedBox.lastSeen = this.frameCount;
+            newBoxes.push(cachedBox);
+            matched = true;
+            break;
+          }
+        }
+
+        if (!matched) {
+          newBoxes.push({
+            bbox: { x1, y1, x2, y2 },
+            age: 0,
+            lastSeen: this.frameCount,
+            id: this.nextId++,
+          });
+        }
+      }
+    }
+
+    this.cachedBoxes = newBoxes;
+  }
+
+  private calculateIoU(box1: number[], box2: number[]): number {
+    const x1 = Math.max(box1[0], box2[0]);
+    const y1 = Math.max(box1[1], box2[1]);
+    const x2 = Math.min(box1[2], box2[2]);
+    const y2 = Math.min(box1[3], box2[3]);
+
+    const interWidth = Math.max(0, x2 - x1);
+    const interHeight = Math.max(0, y2 - y1);
+    const interArea = interWidth * interHeight;
+
+    const box1Area = (box1[2] - box1[0]) * (box1[3] - box1[1]);
+    const box2Area = (box2[2] - box2[0]) * (box2[3] - box2[1]);
+
+    const unionArea = box1Area + box2Area - interArea;
+    return unionArea > 0 ? interArea / unionArea : 0;
+  }
+
+  private cleanupCachedBoxes(): void {
+    // Remove boxes that haven't been seen for too long
+    const maxAge = this.detFrequency * 3;
+    this.cachedBoxes = this.cachedBoxes.filter(
+      (box) => this.frameCount - box.lastSeen < maxAge
+    );
+  }
+}

package/src/solution/wholebody.ts

@@ -0,0 +1,130 @@
+/**
+ * Wholebody solution - combines detection and pose estimation
+ * Based on rtmlib Wholebody class
+ * Supports YOLO12 and YOLOX detectors with RTMW pose model
+ */
+
+import { YOLOX } from '../models/yolox';
+import { YOLO12 } from '../models/yolo12';
+import { RTMPose } from '../models/rtmpose';
+import { BBox, ModeType, ModelConfig } from '../types';
+
+export class Wholebody {
+  private detModel: YOLOX | YOLO12;
+  private poseModel: RTMPose;
+  private detectorType: 'yolox' | 'yolo12';
+
+  private static readonly MODE: Record<ModeType, ModelConfig> = {
+    performance: {
+      det: 'https://download.openmmlab.com/mmpose/v1/projects/rtmposev1/onnx_sdk/yolox_m_8xb8-300e_humanart-c2c7a14a.zip',
+      detInputSize: [640, 640],
+      pose: 'https://download.openmmlab.com/mmpose/v1/projects/rtmw/onnx_sdk/rtmw-dw-x-l_simcc-cocktail14_270e-384x288_20231122.zip',
+      poseInputSize: [384, 288],
+    },
+    lightweight: {
+      det: 'https://download.openmmlab.com/mmpose/v1/projects/rtmposev1/onnx_sdk/yolox_tiny_8xb8-300e_humanart-6f3252f9.zip',
+      detInputSize: [416, 416],
+      pose: 'https://download.openmmlab.com/mmpose/v1/projects/rtmw/onnx_sdk/rtmw-dw-l-m_simcc-cocktail14_270e-256x192_20231122.zip',
+      poseInputSize: [256, 192],
+    },
+    balanced: {
+      det: 'https://download.openmmlab.com/mmpose/v1/projects/rtmposev1/onnx_sdk/yolox_m_8xb8-300e_humanart-c2c7a14a.zip',
+      detInputSize: [640, 640],
+      pose: 'https://download.openmmlab.com/mmpose/v1/projects/rtmw/onnx_sdk/rtmw-dw-x-l_simcc-cocktail14_270e-256x192_20231122.zip',
+      poseInputSize: [256, 192],
+    },
+  };
+
+  constructor(
+    det: string | null = null,
+    detInputSize: [number, number] = [640, 640],
+    pose: string | null = null,
+    poseInputSize: [number, number] = [384, 288],
+    mode: ModeType = 'balanced',
+    toOpenpose: boolean = false,
+    backend: 'onnxruntime' = 'onnxruntime',
+    device: string = 'cpu',
+    detectorType: 'yolox' | 'yolo12' = 'yolox'
+  ) {
+    this.detectorType = detectorType;
+
+    // Use mode config if det/pose not specified
+    let finalDet = det;
+    let finalDetInputSize = detInputSize;
+    let finalPose = pose;
+    let finalPoseInputSize = poseInputSize;
+
+    if (det === null) {
+      finalDet = Wholebody.MODE[mode].det;
+      finalDetInputSize = Wholebody.MODE[mode].detInputSize;
+    }
+
+    if (pose === null) {
+      finalPose = Wholebody.MODE[mode].pose;
+      finalPoseInputSize = Wholebody.MODE[mode].poseInputSize;
+    }
+
+    // Initialize detector based on type
+    if (detectorType === 'yolo12') {
+      this.detModel = new YOLO12(
+        finalDet!,
+        finalDetInputSize,
+        0.45,
+        0.5,
+        backend
+      );
+    } else {
+      this.detModel = new YOLOX(
+        finalDet!,
+        finalDetInputSize,
+        0.45,
+        0.7,
+        backend
+      );
+    }
+
+    this.poseModel = new RTMPose(
+      finalPose!,
+      finalPoseInputSize,
+      toOpenpose,
+      backend
+    );
+  }
+
+  async init(): Promise<void> {
+    await this.detModel.init();
+    await this.poseModel.init();
+  }
+
+  async call(
+    image: Uint8Array,
+    imgWidth: number,
+    imgHeight: number,
+    bboxes?: BBox[]
+  ): Promise<{ keypoints: number[][]; scores: number[] }> {
+    // Run detection if bboxes not provided
+    let finalBboxes = bboxes;
+    if (!finalBboxes || finalBboxes.length === 0) {
+      const detections: any[] = await this.detModel.call(image, imgWidth, imgHeight);
+      // Convert Detection[] or BBox[] to BBox[]
+      finalBboxes = detections.map((d: any) => {
+        if ('bbox' in d) {
+          // Detection type
+          return {
+            x1: d.bbox.x1,
+            y1: d.bbox.y1,
+            x2: d.bbox.x2,
+            y2: d.bbox.y2,
+          };
+        }
+        // Already BBox type
+        return d;
+      });
+    }
+
+    // Run pose estimation
+    const result = await this.poseModel.call(image, imgWidth, imgHeight, finalBboxes);
+
+    return result;
+  }
+}

package/src/solution/wholebody3d.ts

@@ -0,0 +1,125 @@
+/**
+ * Wholebody3D solution - combines detection and 3D pose estimation
+ * Based on rtmlib Wholebody3d class
+ * Uses YOLOX detector with RTMW3D pose model
+ */
+
+import { YOLOX } from '../models/yolox';
+import { RTMPose3D } from '../models/rtmpose3d';
+import { BBox, ModeType, ModelConfig, BackendType } from '../types';
+
+export interface Wholebody3DResult {
+  keypoints: number[][][];
+  scores: number[][];
+  keypointsSimcc: number[][][];
+  keypoints2d: number[][][];
+}
+
+export class Wholebody3D {
+  private detModel: YOLOX;
+  private poseModel: RTMPose3D;
+
+  private static readonly MODE: Record<ModeType, ModelConfig> = {
+    performance: {
+      det: 'https://download.openmmlab.com/mmpose/v1/projects/rtmposev1/onnx_sdk/yolox_m_8xb8-300e_humanart-c2c7a14a.zip',
+      detInputSize: [640, 640],
+      pose: 'https://huggingface.co/Soykaf/RTMW3D-x/resolve/main/onnx/rtmw3d-x_8xb64_cocktail14-384x288-b0a0eab7_20240626.onnx',
+      poseInputSize: [288, 384],  // [width=288, height=384] - creates tensor [1,3,384,288]
+    },
+    lightweight: {
+      det: 'https://download.openmmlab.com/mmpose/v1/projects/rtmposev1/onnx_sdk/yolox_tiny_8xb8-300e_humanart-6f3252f9.zip',
+      detInputSize: [416, 416],
+      pose: 'https://huggingface.co/Soykaf/RTMW3D-x/resolve/main/onnx/rtmw3d-x_8xb64_cocktail14-384x288-b0a0eab7_20240626.onnx',
+      poseInputSize: [192, 256],  // [width=192, height=256]
+    },
+    balanced: {
+      det: 'https://download.openmmlab.com/mmpose/v1/projects/rtmposev1/onnx_sdk/yolox_m_8xb8-300e_humanart-c2c7a14a.zip',
+      detInputSize: [640, 640],
+      pose: 'https://huggingface.co/Soykaf/RTMW3D-x/resolve/main/onnx/rtmw3d-x_8xb64_cocktail14-384x288-b0a0eab7_20240626.onnx',
+      poseInputSize: [288, 384],  // [width=288, height=384] - creates tensor [1,3,384,288]
+    },
+  };
+
+  constructor(
+    det: string | null = null,
+    detInputSize: [number, number] = [640, 640],
+    pose: string | null = null,
+    poseInputSize: [number, number] = [288, 384],  // [width=288, height=384]
+    mode: ModeType = 'balanced',
+    toOpenpose: boolean = false,
+    backend: BackendType = 'webgpu'
+  ) {
+    // Use mode config if det/pose not specified
+    let finalDet = det;
+    let finalDetInputSize = detInputSize;
+    let finalPose = pose;
+    let finalPoseInputSize = poseInputSize;
+
+    if (det === null) {
+      finalDet = Wholebody3D.MODE[mode].det;
+      finalDetInputSize = Wholebody3D.MODE[mode].detInputSize;
+    }
+
+    if (pose === null) {
+      finalPose = Wholebody3D.MODE[mode].pose;
+      finalPoseInputSize = Wholebody3D.MODE[mode].poseInputSize;
+    }
+
+    this.detModel = new YOLOX(
+      finalDet!,
+      finalDetInputSize,
+      0.45,
+      0.7,
+      backend
+    );
+
+    this.poseModel = new RTMPose3D(
+      finalPose!,
+      finalPoseInputSize,
+      toOpenpose,
+      backend
+    );
+  }
+
+  async init(): Promise<void> {
+    await this.detModel.init();
+    await this.poseModel.init();
+  }
+
+  async call(
+    image: Uint8Array,
+    imgWidth: number,
+    imgHeight: number,
+    bboxes?: BBox[]
+  ): Promise<Wholebody3DResult> {
+    // Run detection if bboxes not provided
+    let finalBboxes = bboxes;
+    if (!finalBboxes || finalBboxes.length === 0) {
+      const detections: any[] = await this.detModel.call(image, imgWidth, imgHeight);
+      // Convert Detection[] or BBox[] to BBox[]
+      finalBboxes = detections.map((d: any) => {
+        if ('bbox' in d) {
+          // Detection type
+          return {
+            x1: d.bbox.x1,
+            y1: d.bbox.y1,
+            x2: d.bbox.x2,
+            y2: d.bbox.y2,
+          };
+        }
+        // Already BBox type
+        return d;
+      });
+    }
+
+    // Run 3D pose estimation
+    const result = await this.poseModel.call(
+      image,
+      imgWidth,
+      imgHeight,
+      finalBboxes
+    );
+
+    return result;
+  }
+}

package/src/types/index.ts

@@ -0,0 +1,62 @@
+/**
+ * Basic types for rtmlib-ts
+ * Based on rtmlib Python library
+ */
+
+export interface Keypoint {
+  x: number;
+  y: number;
+  score: number;
+  id: number;
+}
+
+export interface BodyResult {
+  keypoints: Array<Keypoint | null>;
+  totalScore: number;
+  totalParts: number;
+}
+
+export type HandResult = Keypoint[];
+export type FaceResult = Keypoint[];
+
+export interface PoseResult {
+  body: BodyResult;
+  leftHand: HandResult | null;
+  rightHand: HandResult | null;
+  face: FaceResult | null;
+}
+
+export interface BBox {
+  x1: number;
+  y1: number;
+  x2: number;
+  y2: number;
+}
+
+export interface Detection {
+  bbox: BBox;
+  score: number;
+  classId: number;
+}
+
+export interface ModelConfig {
+  det: string;
+  detInputSize: [number, number];
+  pose: string;
+  poseInputSize: [number, number];
+}
+
+export type ModeType = 'performance' | 'lightweight' | 'balanced';
+
+export type BackendType = 'opencv' | 'onnxruntime' | 'openvino' | 'wasm' | 'webgpu';
+export type DeviceType = 'cpu' | 'cuda' | 'mps' | string;
+
+export interface ImageData {
+  data: Uint8Array;
+  width: number;
+  height: number;
+  channels: number;
+}
+
+export type RGBImage = ImageData;
+export type BGRImage = ImageData;