SemanticSegmentationを行うnpmパッケージを作ったので使い方を説明してみる

はじめに

先日、大きな画像を分割してsemantic segmentationを行うnpm パッケージを作成したので、使い方の説明を兼ねてデモの作り方を説明してみます。

このパッケージは、スマホなどの性能の限られたデバイス上で精度と応答速度のトレードオフを調整しながらSemanticSegmentationするためのものです。
今回のデモとしては、このようなものを作ります。バーコードを検出してセグメンテーションを行います。

前回のnpmパッケージを作成した際の投稿はこちらです。
https://qiita.com/wok/items/a6445d17724ed700e448

前準備

今回はReactでデモを作成します。
まずは環境構築をしましょう

$ create-react-app demo --typescript

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

それではvirtual背景のパッケージをインストールします。

$ npm install scalable-semantic-segmentation-js
$ node node_modules/scalable-semantic-segmentation-js/bin/install_worker.js public
file is copied

使用するモデルの準備

次に使用するSemantic Segmentationのモデルを用意します。
このモデルは、[batch, height, width, channels]のshapeのテンソルを入力として想定しています。

$ ls public/WEB_MODEL/300x300_0.10/
group1-shard1of1.bin  model.json

デモのコード

準備ができたら、ソースコードを作成します。
ここでは、重要と思われる部分のみ解説します。ソース全体は下記のリポジトリに置いてあります。

本デモでは、Reactコンポーネントのメンバー変数としてモジュールのクラスのインスタンスを作っておきます。

  scalableSS:ScalableSemanticSegmentation = new ScalableSemanticSegmentation()

では、componentDidMoutからみて行きます。

componentDidMount() {
      console.log('Initializing')

      const initWorkerPromise = this.initWorker()                          // <-- (1)

      if (navigator.mediaDevices && navigator.mediaDevices.getUserMedia) { // <-- (2)
          const webCamPromise = navigator.mediaDevices
              .getUserMedia({
                  audio: false,
                  video: DisplayConstraintOptions[this.state.videoResolution]
              })
              .then(stream => {
                  console.log(this.videoRef)
                  this.videoRef.current!.srcObject = stream;               // <-- (3)
                  return new Promise((resolve, reject) => {
                      this.videoRef.current!.onloadedmetadata = () => {
                          resolve();
                      };
                  });
              });

          Promise.all([initWorkerPromise, webCamPromise])
              .then((res) => {
                  console.log('Camera and model ready!')
              })
              .catch(error => {
                  console.error(error);
              });
      }           
  }

(1)で今回SemanticSegmentationを行うインスタンスをイニシャライズします。このメソッドの内部は後述します。
(2)でビデオデバイスを取得し、(3)でHTMLVideoElementのソースに設定しています。

次に、initWorkerの中身を見ていきます。

async initWorker() {
      // SemanticSegmentation
      this.scalableSS.addInitializedListener(()=>{ // <-- (1-1)
          const props = this.props as any
          this.setState({initialized:true})
          this.requestScanBarcode()                // <-- (1-2)
      })
      this.scalableSS.addMaskPredictedListeners((maskBitmap:ImageBitmap)=>{// <-- (2-1)
          // 再キャプチャ
          this.requestScanBarcode()                                        // <-- (2-2)

      })

      this.scalableSS.init(
            AIConfig.SS_MODEL_PATH, 
            AIConfig.SPLIT_WIDTH, 
            AIConfig.SPLIT_HEIGHT, 
            AIConfig.SPLIT_MARGIN)                                         // <-- (3)
      return
  }

本メソッドでは、SemanticSegmentationを行うインスタンスをイニシャライズしています。
まず(1-1)で、インスタンス内部で使用するモデルのロードなどの初期化が完了したときのコールバック関数を設定します。
コールバック関数の中では(1-2)バーコードスキャンを行う関数を呼び出しています。
(2-1)は、セグメンテーションが完了したときに呼び出されるコールバック関数を設定しています。
このとき受け取るパラメータがセグメンテーションの結果のビットマップイメージになります。
コールバック関数の中でも(2-2)バーコードスキャンを行う関数を呼び出しています。これにより処理をループさせています。
(3)でインスタンスに使用するモデルの情報と、画像を分割する際に使用するマージン（分割後の隣り合う画像間でオーバラップする領域の割合）を指定します。第１引数がモデルのパス、第２、３引数がモデルで使用する幅と高さ、第４引数がマージンです。

最後に、requestScanBarcodeです。

  requestScanBarcode = async () => {
      console.log('requestScanBarcode')
      const video = this.videoRef.current!
      const controller = this.controllerCanvasRef.current!
      controller.width = this.overlayWidth
      controller.height = this.overlayHeight

      const captureCanvas = captureVideoImageToCanvas(video) <--(1)
      if(captureCanvas.width === 0){
          captureCanvas.remove()
          window.requestAnimationFrame(this.requestScanBarcode);
          return
      }
      this.scalableSS.predict(captureCanvas, 
                              this.state.colnum,
                              this.state.rownum) <--(2)
      captureCanvas.remove()
  }

ここでは、SemanticSegmentationを行う対象となる画像を取得し、SemanticSegmentionを実行しています。
(1)で、SemanticSegmentationの対象となる画像をHTMLVideoElementから取得しています。この画像を(2)の引数としてモジュールのインスタンスに与え、SemanticSegmentationを実行させています。第２，３引数は分割する行数と列数です。

以上で、セグメンテーションを行う処理は完了です。

デバッグ用に、セグメンテーションの様子や、グリッドの情報を表示することもできます。

                  <Label basic size="tiny" color={this.state.showSS?"red":"grey"} onClick={()=>{
                      const newValue = !this.state.showSS
                      this.scalableSS.previewCanvas = newValue ? this.workerSSMaskMonitorCanvasRef.current! : null
                      this.setState({showSS:newValue})
                  }}>ss</Label>
                  <Label basic size="tiny" color={this.state.showGrid?"red":"grey"} onClick={()=>{
                      const newValue = !this.state.showGrid
                      this.scalableSS.girdDrawCanvas = newValue ? this.controllerCanvasRef.current! : null
                      this.setState({showGrid:!this.state.showGrid})
                  }}>grid</Label>

this.scalableSS.previewCanvasとthis.scalableSS.girdDrawCanvasにそれぞれ描画するHTMLCanvasElementを設定してください。

デモ

ソースコードとnpmパッケージ

本ソースコードは下記のリポジトリに格納してあります。
https://github.com/FLECT-DEV-TEAM/ScalableSemanticSegmentationjs_demo

npmパッケージのページは次のURLになります。
https://www.npmjs.com/package/scalable-semantic-segmentation-js