【TensorFlow.js】COCO SSDで物体検出を試してみる

　今回は TensorFlow.js の学習済みモデルの MobileNet を使用して物体検出を試してみます。 Angular で書いてますが、 Angular を使う必要はないです。

　コードは前回のコードを引き続き使っていきます。

　完成しているコードは下記にあります。

Angular-sample/tensorflow at main · tsuneken5/Angular-sample

Contribute to tsuneken5/Angular-sample development by creating an account on GitHub.

物体検出とは
環境
パッケージのインストール
サービスの作成
コンポーネントの作成
実行結果

物体検出とは

　物体検出とは、画像から物体の位置、種類、個数を特定する技術のことです。

環境

node: 18.20.0
Angular CLI：17.0.0
TensorFlow.js：4.22.0
TensorFlow-model / COCO SSD：2.2.3

パッケージのインストール

$ npm install @tensorflow/tfjs
$ npm install @tensorflow-models/coco-ssd

　前回で @tensorflow/tfjs はインストールしてますので、不要な方はスキップしてください。

GitHub - tensorflow/tfjs: A WebGL accelerated JavaScript library for training and deploying ML models.

A WebGL accelerated JavaScript library for training and deploying ML models. - tensorflow/tfjs

tfjs-models/coco-ssd at master · tensorflow/tfjs-models

Pretrained models for TensorFlow.js. Contribute to tensorflow/tfjs-models development by creating an account on GitHub.

　また、 tsconfig.json に "skipLibCheck": true を記載して型チェックをスキップさせるのを忘れないでください。

サービスの作成

$ ng generate service service/canvas

  public async uploadImage(imgStr: string, maxWidth: number, canvas: HTMLCanvasElement): Promise<void> {
    const img = new Image();
    const ctx = canvas.getContext('2d') as CanvasRenderingContext2D;
    let scaleRaito: number = 1;

    img.src = imgStr;

    return new Promise((resolve, reject) => {
      img.onload = () => {
        let imageWidth = img.naturalWidth
        if (imageWidth > maxWidth) {
          imageWidth = maxWidth;
          scaleRaito = imageWidth / img.naturalWidth;
        }
        canvas.width = imageWidth;
        let imageHeight = img.naturalHeight * scaleRaito;
        canvas.height = imageHeight;

        resolve(ctx.drawImage(img, 0, 0, imageWidth, imageHeight));
      };
      img.onerror = (error) => {
        console.log(error);
        reject;
      }
    });
  }

  public drawLabel(label: string, x: number, y: number, canvas: HTMLCanvasElement): void {
    const ctx = canvas.getContext('2d') as CanvasRenderingContext2D;

    ctx.save();
    ctx.font = '18px Meiryo';
    ctx.textBaseline = 'top';
    ctx.fillStyle = 'rgba(' + [0, 0, 255, 0.5] + ')';
    let width = ctx.measureText(label).width;
    ctx.fillRect(x, y, width, parseInt(ctx.font, 10));
    ctx.fillStyle = '#fff';
    ctx.fillText(label, x, y);

    ctx.restore();
  }

  public drawBox(box: [number, number, number, number], canvas: HTMLCanvasElement): void {
    const ctx = canvas.getContext('2d') as CanvasRenderingContext2D;
    ctx.strokeStyle = 'blue';

    ctx.strokeRect(...box);
  }

uploadImage

　base64形式の画像を canvas に描画するためのメソッドです。

drawLabel

　canvas に文字を描画するためのメソッドです。このメソッドで推定結果を描画します。

drawBox

　canvas に四角を描画するメソッドです。物体検出のバウンディングボックスを描画するのに使用します。

コンポーネントの作成

$ ng generate component component/page/coco-ssd

import { Component } from '@angular/core';

import * as tf from '@tensorflow/tfjs';
import * as cocoSsd from '@tensorflow-models/coco-ssd';

import { ImageUploaderComponent } from '../../share/image-uploader/image-uploader.component';
import { LoadingSpinnerComponent } from '../../share/loading-spinner/loading-spinner.component';
import { LoadingSpinnerService } from '../../../service/loading-spinner.service';
import { CanvasService } from '../../../service/canvas.service';

@Component({
  selector: 'app-coco-ssd',
  standalone: true,
  imports: [ImageUploaderComponent, LoadingSpinnerComponent],
  templateUrl: './coco-ssd.component.html',
  styleUrl: './coco-ssd.component.css'
})
export class CocoSsdComponent {
  private model!: cocoSsd.ObjectDetection;
  private canvas!: HTMLCanvasElement;

  constructor(
    private loadingSpinnerService: LoadingSpinnerService,
    private canvasService: CanvasService
  ) { }

  private async loadModel(): Promise<void> {
    this.loadingSpinnerService.show();
    await tf.ready();
    this.model = await cocoSsd.load({ base: 'mobilenet_v1' });
    this.loadingSpinnerService.hide();
  }

  private async detect(): Promise<cocoSsd.DetectedObject[]> {
    const predictions = await this.model.detect(this.canvas);
    return predictions;
  }

  public async startDetected(image: string): Promise<void> {
    const parent = this.canvas.parentElement as HTMLElement;
    const width = parent.clientWidth;
    await this.canvasService.uploadImage(image, width, this.canvas);

    const results = await this.detect();
    console.log(results);

    for (const result of results) {
      await this.canvasService.drawBox(result.bbox, this.canvas);
      await this.canvasService.drawLabel(result.class, result.bbox[0], result.bbox[1], this.canvas);
    }
  }

  async ngOnInit() {
    this.loadModel();
  }

  ngAfterViewInit() {
    this.canvas = document.getElementById('myCanvas') as HTMLCanvasElement;
  }
}

loadModel

　モデルをロードします。

base ・・・モデルを指定します。mobilenet_v1 、 mobilenet_v2 、 lite_mobilenet_v2 から選択できます。

detect

　canvas に描画した画像を渡して物体検出を行います。下記のような配列で判定結果が返ってきます。

[
    {
        "bbox": [
            301.82926177978516,
            303.17686265707016,
            698.4287643432617,
            508.2324145436287
        ],
        "class": "horse",
        "score": 0.9725172519683838
    },
    {
        "bbox": [
            610.6569671630859,
            163.37862837314606,
            170.42282104492188,
            426.66063392162323
        ],
        "class": "person",
        "score": 0.6437351703643799
    },
    {
        "bbox": [
            318.3602523803711,
            323.0781354904175,
            73.37371826171875,
            92.26659691333771
        ],
        "class": "person",
        "score": 0.5036835074424744
    }
]

実行結果

元画像

https://pixabay.com/photos/horse-rider-rodeo-cowboy-show-8209533/

実行結果