Class CocoSsdTFJS

tasks/src/tasks/object_detection/cocossd_tfjs.ts:106–148

 * @doc {heading: 'Object Detection', subheading: 'Models'}
 */
export class CocoSsdTFJS extends ObjectDetector<CocoSsdTFJSInferenceOptions> {
  constructor(
      private cocoSsdModel?: cocoSsd.ObjectDetection,
      private loadingOptions?: CocoSsdTFJSLoadingOptions) {
    super();
  }

  async predict(
      img: ImageData|HTMLImageElement|HTMLCanvasElement|HTMLVideoElement,
      infereceOptions?: CocoSsdTFJSInferenceOptions):
      Promise<ObjectDetectionResult> {
    if (!this.cocoSsdModel) {
      throw new Error('source model is not loaded');
    }
    await ensureTFJSBackend(this.loadingOptions);
    const cocoSsdResults = await this.cocoSsdModel.detect(
        img, infereceOptions ? infereceOptions.maxNumBoxes : undefined,
        infereceOptions ? infereceOptions.minScore : undefined);
    const objects: DetectedObject[] = cocoSsdResults.map(result => {
      return {
        boundingBox: {
          originX: result.bbox[0],
          originY: result.bbox[1],
          width: result.bbox[2],
          height: result.bbox[3],
        },
        className: result.class,
        score: result.score,
      };
    });
    const finalResult: ObjectDetectionResult = {
      objects,
    };
    return finalResult;
  }

  cleanUp() {
    if (!this.cocoSsdModel) {
      throw new Error('source model is not loaded');
    }
    return this.cocoSsdModel.dispose();
  }
}

export const cocoSsdTfjsLoader = new CocoSsdTFJSLoader();