67 lines
2.1 KiB
TypeScript
67 lines
2.1 KiB
TypeScript
import type { AIProvider, DetectedObject, MediaItem } from '@photo-gallery/sdk';
|
|
|
|
/**
|
|
* Free, in-browser AI provider using TensorFlow.js + COCO-SSD object detection.
|
|
* No API key, no server — the model (~6 MB) downloads from the public TF model
|
|
* CDN on first use and runs on the GPU via the WebGL backend.
|
|
*
|
|
* Detects 80 COCO classes (person, dog, cat, car, laptop, chair, dining table,
|
|
* tv, bottle, cup, …) which power search, find-similar and the Objects browser.
|
|
*/
|
|
export function createTensorflowProvider(): AIProvider {
|
|
// Loaded lazily and memoized so the heavy bundle/model is fetched only once.
|
|
let modelPromise: Promise<CocoModel> | null = null;
|
|
|
|
async function getModel(): Promise<CocoModel> {
|
|
if (!modelPromise) {
|
|
modelPromise = (async () => {
|
|
const tf = await import('@tensorflow/tfjs');
|
|
try {
|
|
await tf.setBackend('webgl');
|
|
} catch {
|
|
// Fall back to the default backend if WebGL is unavailable.
|
|
}
|
|
await tf.ready();
|
|
const cocoSsd = await import('@tensorflow-models/coco-ssd');
|
|
return cocoSsd.load({ base: 'lite_mobilenet_v2' }) as unknown as CocoModel;
|
|
})();
|
|
}
|
|
return modelPromise;
|
|
}
|
|
|
|
return {
|
|
name: 'tensorflow-coco-ssd',
|
|
|
|
async detectObjects(item: MediaItem, image): Promise<DetectedObject[]> {
|
|
const model = await getModel();
|
|
const el = image as HTMLImageElement;
|
|
const w = el.naturalWidth || el.width || item.width || 1;
|
|
const h = el.naturalHeight || el.height || item.height || 1;
|
|
const predictions = await model.detect(el, 20, 0.4);
|
|
return predictions.map((p) => ({
|
|
label: p.class,
|
|
confidence: p.score,
|
|
box: {
|
|
x: p.bbox[0] / w,
|
|
y: p.bbox[1] / h,
|
|
width: p.bbox[2] / w,
|
|
height: p.bbox[3] / h,
|
|
},
|
|
}));
|
|
},
|
|
};
|
|
}
|
|
|
|
interface CocoPrediction {
|
|
bbox: [number, number, number, number];
|
|
class: string;
|
|
score: number;
|
|
}
|
|
interface CocoModel {
|
|
detect(
|
|
img: HTMLImageElement | HTMLCanvasElement | HTMLVideoElement,
|
|
maxNumBoxes?: number,
|
|
minScore?: number,
|
|
): Promise<CocoPrediction[]>;
|
|
}
|