rtmlib-ts 0.0.2

package/playground/package.json

@@ -0,0 +1,27 @@
+{
+  "name": "playground",
+  "version": "0.1.0",
+  "private": true,
+  "scripts": {
+    "dev": "next dev",
+    "build": "next build",
+    "start": "next start",
+    "lint": "eslint"
+  },
+  "dependencies": {
+    "next": "16.1.6",
+    "react": "19.2.3",
+    "react-dom": "19.2.3",
+    "onnxruntime-web": "^1.23.0"
+  },
+  "devDependencies": {
+    "@tailwindcss/postcss": "^4",
+    "@types/node": "^20",
+    "@types/react": "^19",
+    "@types/react-dom": "^19",
+    "eslint": "^9",
+    "eslint-config-next": "16.1.6",
+    "tailwindcss": "^4",
+    "typescript": "^5"
+  }
+}

package/playground/postcss.config.mjs

@@ -0,0 +1,7 @@
+const config = {
+  plugins: {
+    "@tailwindcss/postcss": {},
+  },
+};
+
+export default config;

package/playground/tsconfig.json

@@ -0,0 +1,34 @@
+{
+  "compilerOptions": {
+    "target": "ES2017",
+    "lib": ["dom", "dom.iterable", "esnext"],
+    "allowJs": true,
+    "skipLibCheck": true,
+    "strict": true,
+    "noEmit": true,
+    "esModuleInterop": true,
+    "module": "esnext",
+    "moduleResolution": "bundler",
+    "resolveJsonModule": true,
+    "isolatedModules": true,
+    "jsx": "react-jsx",
+    "incremental": true,
+    "plugins": [
+      {
+        "name": "next"
+      }
+    ],
+    "paths": {
+      "@/*": ["./*"]
+    }
+  },
+  "include": [
+    "next-env.d.ts",
+    "**/*.ts",
+    "**/*.tsx",
+    ".next/types/**/*.ts",
+    ".next/dev/types/**/*.ts",
+    "**/*.mts"
+  ],
+  "exclude": ["node_modules"]
+}

package/src/core/base.ts

@@ -0,0 +1,66 @@
+/**
+ * BaseTool - Abstract base class for all models
+ * Handles ONNX model loading and inference
+ * Compatible with onnxruntime-web (browser) and onnxruntime-node
+ */
+
+import * as ort from 'onnxruntime-web';
+import { BackendType } from '../types/index.js';
+
+export abstract class BaseTool {
+  protected session: ort.InferenceSession | null = null;
+  protected modelPath: string;
+  protected modelInputSize: [number, number];
+  protected mean: number[] | null;
+  protected std: number[] | null;
+  protected backend: BackendType;
+
+  constructor(
+    modelPath: string,
+    modelInputSize: [number, number],
+    mean: number[] | null = null,
+    std: number[] | null = null,
+    backend: BackendType = 'webgpu'
+  ) {
+    this.modelPath = modelPath;
+    this.modelInputSize = modelInputSize;
+    this.mean = mean;
+    this.std = std;
+    this.backend = backend;
+  }
+
+  protected async init(): Promise<void> {
+    // Configure ONNX Runtime Web - use CDN for WASM files
+    ort.env.wasm.wasmPaths = 'https://cdn.jsdelivr.net/npm/onnxruntime-web@1.23.0/dist/';
+    ort.env.wasm.simd = true;
+    ort.env.wasm.proxy = false;
+
+    // Load model from path/URL
+    this.session = await ort.InferenceSession.create(this.modelPath, {
+      executionProviders: ['wasm'],
+      graphOptimizationLevel: 'all',
+    });
+
+    console.log(`Loaded model: ${this.modelPath}`);
+  }
+
+  protected async inference(img: Float32Array, inputSize?: [number, number]): Promise<any[]> {
+    if (!this.session) {
+      throw new Error('Session not initialized. Call init() first.');
+    }
+
+    const [h, w] = inputSize || this.modelInputSize;
+
+    // Build input tensor (1, 3, H, W)
+    const inputTensor = new (await import('onnxruntime-web')).Tensor('float32', img, [1, 3, h, w]);
+
+    const feeds: Record<string, any> = {};
+    feeds[this.session.inputNames[0]] = inputTensor;
+
+    const results = await this.session.run(feeds);
+
+    return this.session.outputNames.map((name: string) => results[name]);
+  }
+
+  abstract call(...args: unknown[]): Promise<unknown>;
+}

package/src/core/file.ts

@@ -0,0 +1,141 @@
+/**
+ * File utilities for downloading and loading models
+ */
+
+import * as fs from 'fs';
+import * as path from 'path';
+import * as https from 'https';
+import JSZip from 'jszip';
+
+const CACHE_DIR = path.join(process.env.HOME || process.env.USERPROFILE || '.', '.rtmlib', 'models');
+
+export async function downloadCheckpoint(url: string, localPath?: string): Promise<string> {
+  // If local path provided, use it directly
+  if (localPath && fs.existsSync(localPath)) {
+    console.log(`Using local model: ${localPath}`);
+    return localPath;
+  }
+  
+  const fileName = path.basename(url);
+  const cachePath = path.join(CACHE_DIR, fileName.replace('.zip', '.onnx'));
+  
+  if (fs.existsSync(cachePath)) {
+    console.log(`Using cached model: ${cachePath}`);
+    return cachePath;
+  }
+  
+  console.log(`Downloading model from ${url}`);
+  
+  if (!fs.existsSync(CACHE_DIR)) {
+    fs.mkdirSync(CACHE_DIR, { recursive: true });
+  }
+  
+  const tempPath = path.join(CACHE_DIR, fileName);
+  
+  await downloadFile(url, tempPath);
+  
+  if (fileName.endsWith('.zip')) {
+    await extractZip(tempPath, CACHE_DIR);
+    fs.unlinkSync(tempPath);
+  }
+  
+  return cachePath;
+}
+
+async function downloadFile(url: string, dest: string): Promise<void> {
+  return new Promise((resolve, reject) => {
+    const file = fs.createWriteStream(dest);
+    
+    const download = (url: string) => {
+      https.get(url, (response) => {
+        if (response.statusCode === 302 || response.statusCode === 301) {
+          download(response.headers.location!);
+          return;
+        }
+        
+        response.pipe(file);
+        file.on('finish', () => {
+          file.close();
+          resolve();
+        });
+      }).on('error', reject);
+    };
+    
+    download(url);
+  });
+}
+
+async function extractZip(zipPath: string, dest: string): Promise<void> {
+  const data = fs.readFileSync(zipPath);
+  const zip = await JSZip.loadAsync(data);
+  
+  // Find .onnx file in zip
+  for (const [filename, file] of Object.entries(zip.files)) {
+    if (filename.endsWith('.onnx')) {
+      const content = await file.async('nodebuffer');
+      const onnxPath = path.join(dest, filename);
+      
+      // Create directory if needed
+      const dir = path.dirname(onnxPath);
+      if (!fs.existsSync(dir)) {
+        fs.mkdirSync(dir, { recursive: true });
+      }
+      
+      fs.writeFileSync(onnxPath, content);
+      console.log(`Extracted: ${filename}`);
+      return;
+    }
+  }
+  
+  throw new Error('No .onnx file found in zip');
+}
+
+export function fileExists(filePath: string): boolean {
+  return fs.existsSync(filePath);
+}
+
+export function resolveModelPath(modelPath: string): string {
+  if (modelPath.startsWith('http://') || modelPath.startsWith('https://')) {
+    return modelPath;
+  }
+  return path.resolve(modelPath);
+}
+
+/**
+ * Extract local zip file and return onnx path
+ */
+export async function extractLocalZip(zipPath: string): Promise<string> {
+  if (!fs.existsSync(zipPath)) {
+    throw new Error(`Zip file not found: ${zipPath}`);
+  }
+  
+  const destDir = path.dirname(zipPath);
+  const data = fs.readFileSync(zipPath);
+  const zip = await JSZip.loadAsync(data);
+  
+  // Find .onnx file in zip
+  for (const [filename, file] of Object.entries(zip.files)) {
+    if (filename.endsWith('.onnx')) {
+      const onnxPath = path.join(destDir, filename);
+      
+      if (fs.existsSync(onnxPath)) {
+        console.log(`Using extracted model: ${onnxPath}`);
+        return onnxPath;
+      }
+      
+      const content = await file.async('nodebuffer');
+      
+      // Create directory if needed
+      const dir = path.dirname(onnxPath);
+      if (!fs.existsSync(dir)) {
+        fs.mkdirSync(dir, { recursive: true });
+      }
+      
+      fs.writeFileSync(onnxPath, content);
+      console.log(`Extracted: ${filename} -> ${onnxPath}`);
+      return onnxPath;
+    }
+  }
+  
+  throw new Error('No .onnx file found in zip');
+}

package/src/core/modelCache.ts

@@ -0,0 +1,189 @@
+/**
+ * Model caching utility using Cache API
+ * Caches ONNX models in browser to avoid repeated downloads
+ */
+
+const CACHE_NAME = 'rtmlib-ts-models-v1';
+
+/**
+ * Check if model is available in cache
+ */
+export async function isModelCached(url: string): Promise<boolean> {
+  if (typeof caches === 'undefined') {
+    // Cache API not available (e.g., Node.js)
+    return false;
+  }
+  
+  try {
+    const cache = await caches.open(CACHE_NAME);
+    const response = await cache.match(url);
+    return !!response;
+  } catch (error) {
+    console.warn(`[ModelCache] Failed to check cache for ${url}:`, error);
+    return false;
+  }
+}
+
+/**
+ * Get model from cache or fetch from network
+ * @param url - Model URL
+ * @param forceRefresh - Force refresh from network
+ */
+export async function getCachedModel(url: string, forceRefresh: boolean = false): Promise<ArrayBuffer> {
+  if (typeof caches === 'undefined') {
+    // Cache API not available, fetch directly
+    console.log(`[ModelCache] Cache API not available, fetching from network`);
+    return fetchModelFromNetwork(url);
+  }
+  
+  try {
+    const cache = await caches.open(CACHE_NAME);
+    
+    // Try to get from cache first
+    if (!forceRefresh) {
+      const cachedResponse = await cache.match(url);
+      if (cachedResponse) {
+        console.log(`[ModelCache] ✅ Hit for ${url}`);
+        return await cachedResponse.arrayBuffer();
+      }
+      console.log(`[ModelCache] ❌ Miss for ${url}, fetching from network...`);
+    }
+    
+    // Fetch from network
+    const networkResponse = await fetchModelFromNetwork(url);
+    
+    // Cache the response
+    const responseToCache = new Response(networkResponse, {
+      headers: {
+        'Content-Type': 'application/octet-stream',
+      },
+    });
+    
+    await cache.put(url, responseToCache);
+    console.log(`[ModelCache] 💾 Cached ${url}`);
+    
+    return networkResponse;
+  } catch (error) {
+    console.error(`[ModelCache] Failed to get/cache model ${url}:`, error);
+    throw error;
+  }
+}
+
+/**
+ * Fetch model from network with progress tracking
+ */
+async function fetchModelFromNetwork(url: string): Promise<ArrayBuffer> {
+  const response = await fetch(url);
+  
+  if (!response.ok) {
+    throw new Error(`Failed to fetch model: HTTP ${response.status} ${response.statusText}`);
+  }
+  
+  return await response.arrayBuffer();
+}
+
+/**
+ * Preload and cache multiple models
+ */
+export async function preloadModels(urls: string[]): Promise<void> {
+  console.log(`[ModelCache] Preloading ${urls.length} model(s)...`);
+  
+  const results = await Promise.allSettled(
+    urls.map(url => getCachedModel(url))
+  );
+  
+  const success = results.filter(r => r.status === 'fulfilled').length;
+  const failed = results.filter(r => r.status === 'rejected').length;
+  
+  console.log(`[ModelCache] Preload complete: ${success} succeeded, ${failed} failed`);
+  
+  results.forEach((result, index) => {
+    if (result.status === 'rejected') {
+      console.error(`[ModelCache] Failed to preload ${urls[index]}:`, result.reason);
+    }
+  });
+}
+
+/**
+ * Clear all cached models
+ */
+export async function clearModelCache(): Promise<void> {
+  if (typeof caches === 'undefined') {
+    return;
+  }
+  
+  try {
+    await caches.delete(CACHE_NAME);
+    console.log('[ModelCache] Cache cleared');
+  } catch (error) {
+    console.error('[ModelCache] Failed to clear cache:', error);
+  }
+}
+
+/**
+ * Get cache size in bytes
+ */
+export async function getCacheSize(): Promise<number> {
+  if (typeof caches === 'undefined' || !navigator.storage) {
+    return 0;
+  }
+  
+  try {
+    const cache = await caches.open(CACHE_NAME);
+    const keys = await cache.keys();
+    let totalSize = 0;
+    
+    for (const request of keys) {
+      const response = await cache.match(request);
+      if (response) {
+        const blob = await response.blob();
+        totalSize += blob.size;
+      }
+    }
+    
+    return totalSize;
+  } catch (error) {
+    console.warn('[ModelCache] Failed to get cache size:', error);
+    return 0;
+  }
+}
+
+/**
+ * Get cache info
+ */
+export async function getCacheInfo(): Promise<{
+  cachedModels: string[];
+  totalSize: number;
+  totalSizeFormatted: string;
+}> {
+  if (typeof caches === 'undefined') {
+    return { cachedModels: [], totalSize: 0, totalSizeFormatted: '0 B' };
+  }
+  
+  try {
+    const cache = await caches.open(CACHE_NAME);
+    const keys = await cache.keys();
+    const cachedModels = keys.map(k => k.url);
+    const totalSize = await getCacheSize();
+    
+    return {
+      cachedModels,
+      totalSize,
+      totalSizeFormatted: formatBytes(totalSize),
+    };
+  } catch (error) {
+    console.warn('[ModelCache] Failed to get cache info:', error);
+    return { cachedModels: [], totalSize: 0, totalSizeFormatted: '0 B' };
+  }
+}
+
+/**
+ * Format bytes to human-readable string
+ */
+function formatBytes(bytes: number): string {
+  if (bytes === 0) return '0 B';
+  const k = 1024;
+  const sizes = ['B', 'KB', 'MB', 'GB'];
+  const i = Math.floor(Math.log(bytes) / Math.log(k));
+  return `${(bytes / Math.pow(k, i)).toFixed(2)} ${sizes[i]}`;
+}

package/src/core/posePostprocessing.ts

@@ -0,0 +1,91 @@
+/**
+ * Post-processing utilities for pose estimation
+ */
+
+export function getSimccMaximum(
+  simccX: Float32Array,
+  simccY: Float32Array
+): { locations: number[]; scores: number[] } {
+  const numKeypoints = simccX.length / 2; // Assuming split_ratio = 2
+  
+  const locations: number[] = [];
+  const scores: number[] = [];
+  
+  for (let i = 0; i < numKeypoints; i++) {
+    // Find argmax for x
+    let maxX = -Infinity;
+    let argmaxX = 0;
+    for (let j = 0; j < 2; j++) {
+      const val = simccX[i * 2 + j];
+      if (val > maxX) {
+        maxX = val;
+        argmaxX = j;
+      }
+    }
+    
+    // Find argmax for y
+    let maxY = -Infinity;
+    let argmaxY = 0;
+    for (let j = 0; j < 2; j++) {
+      const val = simccY[i * 2 + j];
+      if (val > maxY) {
+        maxY = val;
+        argmaxY = j;
+      }
+    }
+    
+    locations.push(argmaxX, argmaxY);
+    scores.push((maxX + maxY) / 2);
+  }
+  
+  return { locations, scores };
+}
+
+export function 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, // right_ear
+    17, // left_ear
+  ];
+  
+  const openposeKeypoints: number[][] = [];
+  const openposeScores: number[] = [];
+  
+  for (let i = 0; i < 18; 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 };
+}

package/src/core/postprocessing.ts

@@ -0,0 +1,93 @@
+/**
+ * Post-processing utilities for object detection
+ */
+
+import { BBox } from '../types/index.js';
+
+export function multiclassNms(
+  boxes: number[][],
+  scores: number[][],
+  nmsThr: number,
+  scoreThr: number
+): { boxes: number[][]; scores: number[]; classIds: number[] } | null {
+  const numClasses = scores[0].length;
+  const numBoxes = boxes.length;
+  
+  const allDetections: Array<{ box: number[]; score: number; classId: number }> = [];
+  
+  for (let c = 0; c < numClasses; c++) {
+    for (let i = 0; i < numBoxes; i++) {
+      if (scores[i][c] > scoreThr) {
+        allDetections.push({
+          box: boxes[i],
+          score: scores[i][c],
+          classId: c,
+        });
+      }
+    }
+  }
+  
+  allDetections.sort((a, b) => b.score - a.score);
+  
+  const keep: typeof allDetections = [];
+  
+  while (allDetections.length > 0) {
+    const current = allDetections.shift()!;
+    keep.push(current);
+    
+    for (let i = allDetections.length - 1; i >= 0; i--) {
+      const other = allDetections[i];
+      if (current.classId !== other.classId) continue;
+      
+      const iou = calculateIoU(current.box, other.box);
+      if (iou > nmsThr) {
+        allDetections.splice(i, 1);
+      }
+    }
+  }
+  
+  if (keep.length === 0) return null;
+  
+  return {
+    boxes: keep.map(d => d.box),
+    scores: keep.map(d => d.score),
+    classIds: keep.map(d => d.classId),
+  };
+}
+
+function 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;
+}
+
+export function nms(boxes: number[][], scores: number[], nmsThr: number): number[] {
+  const indices: number[] = [];
+  const sortedIndices = scores.map((s, i) => i).sort((a, b) => scores[b] - scores[a]);
+  
+  while (sortedIndices.length > 0) {
+    const current = sortedIndices.shift()!;
+    indices.push(current);
+    
+    for (let i = sortedIndices.length - 1; i >= 0; i--) {
+      const other = sortedIndices[i];
+      const iou = calculateIoU(boxes[current], boxes[other]);
+      if (iou > nmsThr) {
+        sortedIndices.splice(i, 1);
+      }
+    }
+  }
+  
+  return indices;
+}