rtmlib-ts 0.0.2

package/src/core/preprocessing.ts

@@ -0,0 +1,127 @@
+/**
+ * Pre-processing utilities for pose estimation
+ */
+
+export function bboxXyxy2cs(
+  bbox: [number, number, number, number],
+  padding: number = 1.25
+): { center: [number, number]; scale: [number, number] } {
+  const [x1, y1, x2, y2] = bbox;
+
+  const center: [number, number] = [(x1 + x2) / 2, (y1 + y2) / 2];
+
+  const w = x2 - x1;
+  const h = y2 - y1;
+
+  // Python: scale = w * padding, h * padding (different values!)
+  const scale: [number, number] = [w * padding, h * padding];
+
+  return { center, scale };
+}
+
+export function topDownAffine(
+  imageSize: [number, number],
+  scale: [number, number],
+  center: [number, number],
+  img: Uint8Array,
+  imgWidth: number,
+  imgHeight: number
+): { resizedImg: Float32Array; scale: [number, number] } {
+  const [w, h] = imageSize;
+
+  const srcW = scale[0];
+  const srcH = scale[1];
+
+  // Calculate transformation matrix
+  const scaleX = w / srcW;
+  const scaleY = h / srcH;
+
+  // Create output array
+  const outputSize = w * h * 3;
+  const resizedImg = new Float32Array(outputSize);
+
+  // Simple bilinear interpolation
+  for (let y = 0; y < h; y++) {
+    for (let x = 0; x < w; x++) {
+      // Map output coordinates to input coordinates
+      const srcX = (x / w) * srcW + (center[0] - srcW / 2);
+      const srcY = (y / h) * srcH + (center[1] - srcH / 2);
+
+      // Get the four nearest pixels
+      const x0 = Math.floor(srcX);
+      const y0 = Math.floor(srcY);
+      const x1 = x0 + 1;
+      const y1 = y0 + 1;
+
+      const dx = srcX - x0;
+      const dy = srcY - y0;
+
+      // Sample from input image with bounds checking
+      for (let c = 0; c < 3; c++) {
+        const p00 = getPixel(img, imgWidth, imgHeight, x0, y0, c);
+        const p10 = getPixel(img, imgWidth, imgHeight, x1, y0, c);
+        const p01 = getPixel(img, imgWidth, imgHeight, x0, y1, c);
+        const p11 = getPixel(img, imgWidth, imgHeight, x1, y1, c);
+
+        // Bilinear interpolation
+        const value = p00 * (1 - dx) * (1 - dy) +
+                      p10 * dx * (1 - dy) +
+                      p01 * (1 - dx) * dy +
+                      p11 * dx * dy;
+
+        resizedImg[y * w * 3 + x * 3 + c] = value;
+      }
+    }
+  }
+
+  // Return original scale (not model dimensions) for postprocess
+  return { resizedImg, scale };
+}
+
+function getPixel(
+  img: Uint8Array,
+  width: number,
+  height: number,
+  x: number,
+  y: number,
+  channel: number
+): number {
+  if (x < 0 || x >= width || y < 0 || y >= height) {
+    return 0;
+  }
+  return img[y * width * 3 + x * 3 + channel];
+}
+
+export function normalizeImage(
+  img: Float32Array,
+  mean: number[],
+  std: number[]
+): Float32Array {
+  const normalized = new Float32Array(img.length);
+  
+  for (let i = 0; i < img.length; i++) {
+    const channel = i % 3;
+    normalized[i] = (img[i] - mean[channel]) / std[channel];
+  }
+  
+  return normalized;
+}
+
+export function transposeImage(
+  img: Float32Array,
+  height: number,
+  width: number
+): Float32Array {
+  // HWC to CHW
+  const transposed = new Float32Array(img.length);
+  
+  for (let c = 0; c < 3; c++) {
+    for (let h = 0; h < height; h++) {
+      for (let w = 0; w < width; w++) {
+        transposed[c * height * width + h * width + w] = img[h * width * 3 + w * 3 + c];
+      }
+    }
+  }
+  
+  return transposed;
+}

package/src/index.ts

@@ -0,0 +1,69 @@
+/**
+ * rtmlib-ts - Real-Time Multi-Person Pose Estimation Library
+ *
+ * TypeScript port of rtmlib Python library
+ * Based on RTMPose, DWPose, RTMO, RTMW models
+ */
+
+// Models
+export { YOLOX } from './models/yolox';
+export { YOLO12 } from './models/yolo12';
+export { RTMPose } from './models/rtmpose';
+export { RTMPose3D } from './models/rtmpose3d';
+
+// Solutions (High-level APIs)
+export { ObjectDetector, COCO_CLASSES } from './solution/objectDetector';
+export { PoseDetector } from './solution/poseDetector';
+export { Pose3DDetector, type Person3D, type Pose3DStats } from './solution/pose3dDetector';
+export { CustomDetector, type CustomDetectorConfig, type DetectionResult } from './solution/customDetector';
+export { AnimalDetector, ANIMAL_CLASSES, VITPOSE_MODELS, type VitPoseModelType, type DetectedAnimal, type AnimalKeypoint } from './solution/animalDetector';
+export { Wholebody } from './solution/wholebody';
+export { Wholebody3D, type Wholebody3DResult } from './solution/wholebody3d';
+export { Body } from './solution/body';
+export { Hand } from './solution/hand';
+export { BodyWithFeet } from './solution/bodyWithFeet';
+export { PoseTracker } from './solution/poseTracker';
+
+// Visualization
+export { drawBbox, drawSkeleton, drawDetectionsOnCanvas, drawPoseOnCanvas, drawResultsOnCanvas } from './visualization/draw';
+
+// Model caching utilities
+export {
+  getCachedModel,
+  isModelCached,
+  preloadModels,
+  clearModelCache,
+  getCacheSize,
+  getCacheInfo,
+} from './core/modelCache';
+
+// Types
+export type {
+  Keypoint,
+  BodyResult,
+  HandResult,
+  FaceResult,
+  PoseResult,
+  BBox,
+  Detection,
+  ModelConfig,
+  ModeType,
+  BackendType,
+  DeviceType,
+  ImageData,
+  RGBImage,
+  BGRImage,
+} from './types/index';
+
+// Skeleton configurations
+export {
+  coco17,
+  coco133,
+  hand21,
+  halpe26,
+  openpose18,
+  openpose134,
+} from './visualization/skeleton/index';
+
+// Version
+export const VERSION = '0.0.1';

package/src/models/rtmpose.ts

@@ -0,0 +1,265 @@
+/**
+ * RTMPose model for pose estimation
+ * Supports RTMPose, DWPose, RTMW variants
+ * Uses onnxruntime-web for browser compatibility
+ */
+
+import { BaseTool } from '../core/base';
+import { BBox, BackendType } from '../types/index';
+
+export class RTMPose extends BaseTool {
+  private toOpenpose: boolean;
+  private simccSplitRatio: number = 2.0;
+  private initialized: boolean = false;
+
+  private readonly defaultMean: number[] = [123.675, 116.28, 103.53];
+  private readonly defaultStd: number[] = [58.395, 57.12, 57.375];
+
+  constructor(
+    onnxModel: string,
+    modelInputSize: [number, number] = [384, 288],  // [height, width]
+    toOpenpose: boolean = false,
+    backend: BackendType = 'webgpu'
+  ) {
+    super(onnxModel, modelInputSize, null, null, backend);
+    this.toOpenpose = toOpenpose;
+  }
+
+  async init(): Promise<void> {
+    // Web version - model path is direct URL
+    await super.init();
+    this.initialized = true;
+  }
+
+  async call(
+    image: Uint8Array,
+    imgWidth: number,
+    imgHeight: number,
+    bboxes: BBox[] = []
+  ): Promise<{ keypoints: number[][]; scores: number[] }> {
+    if (!this.initialized) {
+      await this.init();
+    }
+
+    if (bboxes.length === 0) {
+      bboxes = [{ x1: 0, y1: 0, x2: imgWidth, y2: imgHeight }];
+    }
+
+    const allKeypoints: number[][][] = [];
+    const allScores: number[][] = [];
+
+    for (const bbox of bboxes) {
+      const { tensor, center, scale, inputSize } = this.preprocess(
+        image,
+        imgWidth,
+        imgHeight,
+        bbox
+      );
+
+      const outputs = await this.inference(tensor, inputSize);
+      const { keypoints, scores } = this.postprocess(outputs[0].data as Float32Array, outputs[1].data as Float32Array, outputs[0].dims, outputs[1].dims, center, scale);
+
+      allKeypoints.push(keypoints);
+      allScores.push(scores);
+    }
+
+    // Flatten results
+    const keypoints = allKeypoints.flat();
+    const scores = allScores.flat();
+
+    if (this.toOpenpose) {
+      const converted = this.convertCocoToOpenpose(keypoints, scores);
+      return converted;
+    }
+
+    return { keypoints, scores };
+  }
+
+  private preprocess(
+    img: Uint8Array,
+    imgWidth: number,
+    imgHeight: number,
+    bbox: BBox
+  ): { tensor: Float32Array; center: [number, number]; scale: [number, number]; inputSize: [number, number] } {
+    const [inputH, inputW] = this.modelInputSize;  // H=384, W=288
+    
+    // Center and scale from bbox with padding
+    const center: [number, number] = [
+      bbox.x1 + (bbox.x2 - bbox.x1) / 2,
+      bbox.y1 + (bbox.y2 - bbox.y1) / 2,
+    ];
+    
+    const bboxWidth = bbox.x2 - bbox.x1;
+    const bboxHeight = bbox.y2 - bbox.y1;
+    const bboxAspectRatio = bboxWidth / bboxHeight;
+    const modelAspectRatio = inputW / inputH;
+    
+    let scaleW: number, scaleH: number;
+    if (bboxAspectRatio > modelAspectRatio) {
+      scaleW = bboxWidth * 1.25;
+      scaleH = scaleW / modelAspectRatio;
+    } else {
+      scaleH = bboxHeight * 1.25;
+      scaleW = scaleH * modelAspectRatio;
+    }
+    
+    const scale: [number, number] = [scaleW, scaleH];
+
+    // Create canvas for cropping
+    const canvas = document.createElement('canvas');
+    const ctx = canvas.getContext('2d')!;
+    canvas.width = inputW;
+    canvas.height = inputH;
+
+    ctx.fillStyle = '#FFFFFF';
+    ctx.fillRect(0, 0, inputW, inputH);
+
+    // Create source canvas from image data
+    const srcCanvas = document.createElement('canvas');
+    const srcCtx = srcCanvas.getContext('2d')!;
+    srcCanvas.width = imgWidth;
+    srcCanvas.height = imgHeight;
+    
+    const srcImageData = srcCtx.createImageData(imgWidth, imgHeight);
+    srcImageData.data.set(img);
+    srcCtx.putImageData(srcImageData, 0, 0);
+
+    // Calculate source region
+    const srcX = center[0] - scaleW / 2;
+    const srcY = center[1] - scaleH / 2;
+
+    // Draw cropped and scaled region
+    ctx.drawImage(
+      srcCanvas,
+      srcX,
+      srcY,
+      scaleW,
+      scaleH,
+      0,
+      0,
+      inputW,
+      inputH
+    );
+
+    const imageData = ctx.getImageData(0, 0, inputW, inputH);
+
+    // Normalize with mean/std
+    const data = new Float32Array(inputW * inputH * 3);
+    for (let i = 0; i < imageData.data.length; i += 4) {
+      const pixelIndex = i / 4;
+      for (let c = 0; c < 3; c++) {
+        const value = imageData.data[i + c];
+        data[c * inputW * inputH + pixelIndex] = (value - this.defaultMean[c]) / this.defaultStd[c];
+      }
+    }
+
+    return {
+      tensor: data,
+      center,
+      scale,
+      inputSize: [inputH, inputW],
+    };
+  }
+
+  private postprocess(
+    simccX: Float32Array,
+    simccY: Float32Array,
+    outputShapeX: number[],
+    outputShapeY: number[],
+    center: [number, number],
+    scale: [number, number]
+  ): { keypoints: number[][]; scores: number[] } {
+    const numKeypoints = outputShapeX[1];
+    const wx = outputShapeX[2];
+    const wy = outputShapeY[2];
+
+    const keypoints: number[][] = [];
+    const scores: number[] = [];
+
+    for (let k = 0; k < numKeypoints; k++) {
+      // Find argmax for x
+      let maxX = -Infinity;
+      let argmaxX = 0;
+      for (let i = 0; i < wx; i++) {
+        const val = simccX[k * wx + i];
+        if (val > maxX) {
+          maxX = val;
+          argmaxX = i;
+        }
+      }
+
+      // Find argmax for y
+      let maxY = -Infinity;
+      let argmaxY = 0;
+      for (let i = 0; i < wy; i++) {
+        const val = simccY[k * wy + i];
+        if (val > maxY) {
+          maxY = val;
+          argmaxY = i;
+        }
+      }
+
+      const score = 0.5 * (maxX + maxY);
+
+      // Normalize to [0, 1] and transform to original image coordinates
+      const normX = argmaxX / wx;
+      const normY = argmaxY / wy;
+
+      const kptX = (normX - 0.5) * scale[0] + center[0];
+      const kptY = (normY - 0.5) * scale[1] + center[1];
+
+      keypoints.push([kptX, kptY]);
+      scores.push(score);
+    }
+
+    return { keypoints, scores };
+  }
+
+  private convertCocoToOpenpose(
+    keypoints: number[][],
+    scores: number[]
+  ): { keypoints: number[][]; scores: number[] } {
+    // COCO 17 keypoints to OpenPose 18 keypoints mapping
+    const cocoToOpenpose: number[] = [
+      0,  // nose
+      1,  // neck (average of shoulders)
+      2,  // right_shoulder
+      3,  // right_elbow
+      4,  // right_wrist
+      5,  // left_shoulder
+      6,  // left_elbow
+      7,  // left_wrist
+      8,  // right_hip
+      9,  // right_knee
+      10, // right_ankle
+      11, // left_hip
+      12, // left_knee
+      13, // left_ankle
+      14, // right_eye
+      15, // left_eye
+      16, // left_ear
+    ];
+
+    const openposeKeypoints: number[][] = [];
+    const openposeScores: number[] = [];
+
+    for (let i = 0; i < 17; i++) {
+      if (i === 1) {
+        // Neck is average of shoulders
+        const rightShoulder = keypoints[2];
+        const leftShoulder = keypoints[5];
+        openposeKeypoints.push([
+          (rightShoulder[0] + leftShoulder[0]) / 2,
+          (rightShoulder[1] + leftShoulder[1]) / 2,
+        ]);
+        openposeScores.push((scores[2] + scores[5]) / 2);
+      } else {
+        const cocoIdx = cocoToOpenpose[i];
+        openposeKeypoints.push([...keypoints[cocoIdx]]);
+        openposeScores.push(scores[cocoIdx]);
+      }
+    }
+
+    return { keypoints: openposeKeypoints, scores: openposeScores };
+  }
+}