Spaces:
Running
Running
| import { pipeline } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@4.2.0'; | |
| export const MODEL_OPTIONS = [ | |
| { | |
| key: 'rtdetr', | |
| label: 'RT-DETR R18', | |
| modelId: 'onnx-community/rtdetr_r18vd', | |
| note: 'Recommended classroom default: modern real-time transformer-style detector.', | |
| dtype: 'q8', | |
| }, | |
| { | |
| key: 'detr', | |
| label: 'DETR ResNet-50', | |
| modelId: 'Xenova/detr-resnet-50', | |
| note: 'Representative DETR-family detector. Larger download; useful for comparison.', | |
| dtype: 'q8', | |
| }, | |
| { | |
| key: 'yolos', | |
| label: 'YOLOS Tiny', | |
| modelId: 'Xenova/yolos-tiny', | |
| note: 'Small browser-friendly detector. Useful for showing that model choice changes results.', | |
| dtype: 'q8', | |
| }, | |
| ]; | |
| const detectorCache = new Map(); | |
| export function getModelOption(key) { | |
| return MODEL_OPTIONS.find((model) => model.key === key) || MODEL_OPTIONS[0]; | |
| } | |
| export async function loadDetector(modelKey, progressCallback) { | |
| const option = getModelOption(modelKey); | |
| if (detectorCache.has(option.key)) return detectorCache.get(option.key); | |
| const detector = await pipeline('object-detection', option.modelId, { | |
| dtype: option.dtype, | |
| progress_callback: progressCallback, | |
| }); | |
| detectorCache.set(option.key, detector); | |
| return detector; | |
| } | |
| export async function detectObjects(imageDataUrl, options = {}) { | |
| const modelKey = options.modelKey || MODEL_OPTIONS[0].key; | |
| const threshold = Number.isFinite(options.threshold) ? options.threshold : 0.5; | |
| const detector = await loadDetector(modelKey, options.progressCallback); | |
| return detector(imageDataUrl, { threshold }); | |
| } | |