yolo-tfjs-vision
Version:
A TypeScript-based YOLO object detection library using TensorFlow.js
174 lines (173 loc) • 6.5 kB
TypeScript
import * as tf from "@tensorflow/tfjs";
import "@tensorflow/tfjs-backend-webgl";
/**
* Interface representing the configuration for the YOLO object detection.
*/
export interface YOLOConfig {
/**
* The URL pointing to the TensorFlow.js YOLO model JSON file.
* This model should be compatible with TensorFlow.js for in-browser execution.
*
* @example "https://example.com/model/model.json"
*/
modelUrl: string;
/**
* Optional array of labels for object detection.
* If omitted, the model will default to using COCO dataset labels.
*
* @example ["person", "car", "dog"]
*/
labels: string[];
/**
* Optional array of colors for detected object labels.
* The color mapping corresponds to the label array.
* If omitted, a default color scheme is used.
*
* @example ["#FF0000", "#00FF00", "#0000FF"]
*/
colors: string[];
/**
* Optional set of labels to filter displayed objects.
* If null, all detected objects will be displayed.
*
* @example new Set(["person", "dog"])
*/
displayLabels: Set<string> | null;
/**
* The minimum confidence score threshold for displaying detected objects.
* Objects with a lower score will be ignored.
*
* @default 0.5
* @example 0.3
*/
scoreThreshold: number;
/**
* The width of the bounding box stroke for detected objects.
* Adjusting this value controls how thick the detection box appears.
*
* @default 2
* @example 10
*/
boxLineWidth: number;
/**
* Determines whether labels should be displayed inside the bounding boxes.
* If `true`, the detected object’s label will be rendered inside its corresponding box.
*
* @default true
* @example false
*/
boxLabels: boolean;
}
/**
* Interface representing the loaded YOLO model.
*/
export interface YOLOModel {
net: tf.GraphModel;
inputShape: number[];
}
/**
* Interface representing the detection result from the YOLO model.
*/
export interface Detection {
/**
* Flattened array of bounding box coordinates.
* Each detection's box is represented by four consecutive numbers in the order [y1, x1, y2, x2].
*
* @example [50, 30, 200, 180, ... ]
*/
boxes: number[];
/**
* Array of confidence scores for each detection.
* Each score indicates the probability that the detection is accurate.
*
* @example [0.95, 0.87, ... ]
*/
scores: number[];
/**
* Array of class indices for each detection.
* Each index corresponds to a specific label defined in the YOLO configuration.
*
* @example [0, 3, ... ]
*/
classes: number[];
/**
* Array of label strings corresponding to the detected classes.
* These provide a human-readable description for each detected object.
*
* @example ["person", "dog", ... ]
*/
labels: string[];
}
/**
* YOLO class for object detection using TensorFlow.js.
*/
declare class YOLO {
/**
* Default configuration for the YOLO model.
*/
config: YOLOConfig;
/**
* Sets up the YOLO configuration.
*
* @param options - Partial configuration options to override the defaults.
* `displayLabels` can be provided as an array of strings.
*/
setup(options: Partial<Omit<YOLOConfig, "displayLabels">> & {
displayLabels?: string[];
}): void;
/**
* Loads the YOLO model from the specified URL.
*
* @returns A promise that resolves to the loaded YOLO model, or null if an error occurs.
*/
loadModel(): Promise<YOLOModel | null>;
/**
* Detects objects in the given source (image, video, or canvas) and draws bounding boxes on the canvas.
*
* @param source - The source element (HTMLImageElement, HTMLVideoElement, or HTMLCanvasElement).
* @param model - The loaded YOLO model.
* @param canvasRef - The canvas element where the detection boxes will be rendered.
* @param callback - A callback function that receives detection data (boxes, scores, classes, labels).
*/
detect(source: HTMLImageElement | HTMLVideoElement | HTMLCanvasElement, model: YOLOModel, canvasRef: HTMLCanvasElement, callback?: (detections: Detection) => void): Promise<void>;
/**
* Continuously detects objects in a video source and updates the canvas in real time.
*
* @param videoSource - The HTMLVideoElement that provides the video feed.
* @param model - The loaded YOLO model.
* @param canvasRef - The canvas element where detection boxes will be rendered.
*/
detectVideo(videoSource: HTMLVideoElement, model: YOLOModel, canvasRef: HTMLCanvasElement): void;
/**
* Preprocesses the source image/video/canvas by resizing it with letterboxing
* to maintain aspect ratio.
*
* @param source - The source element (HTMLImageElement, HTMLVideoElement, or HTMLCanvasElement).
* @param modelWidth - The target width for the model input.
* @param modelHeight - The target height for the model input.
* @returns A tuple containing the preprocessed input tensor and an object with letterbox details:
* - scale: The scaling factor applied.
* - dx: Horizontal padding.
* - dy: Vertical padding.
* - origWidth: Original source width.
* - origHeight: Original source height.
*/
preprocess(source: HTMLImageElement | HTMLVideoElement | HTMLCanvasElement, modelWidth: number, modelHeight: number): [tf.Tensor, {
scale: number;
dx: number;
dy: number;
origWidth: number;
origHeight: number;
}];
/**
* Renders detection boxes and labels on the provided canvas.
*
* @param canvasRef - The canvas element where detections will be drawn.
* @param boxesData - Array of box coordinates (assumed to be in the original image coordinate system).
* @param scoresData - Array of detection confidence scores.
* @param classesData - Array of detected class indices.
* @param ratios - Scaling ratios (typically [1, 1] if boxes are already in original image coordinates).
*/
renderBoxes(canvasRef: HTMLCanvasElement, boxesData: number[] | Float32Array, scoresData: number[] | Float32Array, classesData: number[] | Float32Array, ratios: [number, number]): void;
}
export default YOLO;