公式の CycleGAN チュートリアルをオリジナルデータで動かした時のメモです。Google Colab 上で動かします。

1. データの準備

リンクがある cyclegan.ipynb はデフォルトで tensorflow_datasets の cycle_gan/horse2zebra を利用するようになっています。順当にやるなら tfds new などでカスタムデータセットを作成する形なります。
ただ若干書くことがありそうだったので、以下に書いたように horse2zebra.zip の DL を置換する方法で動かしました。

1.1. horse2zebra.zip の DL を置換

cycle_gan/horse2zebra の実装は大学のサイトから horse2zebra.zip を落としてくる実装になっています。
ただ試すだけであれば、horse2zebra.zip を落としてくる部分を置換すれば良いので、まずデータを以下のように作成します。

• horse2zebra
  • testA
    • testA_0001.jpg
    • testA_0002.jpg
    • ...
  • testB
    • testB_0001.jpg
    • testB_0002.jpg
    • ...
  • trainA
    • trainA_0001.jpg
    • trainA_0002.jpg
    • ...
  • trainB
    • trainB_0001.jpg
    • trainB_0002.jpg
    • ...

上記の構成 (horse2zebra 含め) を zip を固めて Google Drive に上げておき(例: my_horse2zebra.zip)、以下のコードを cyclegan.ipynb の冒頭にいれると、ダウンロードの代わりに zip を利用できます。

## Google Drive をマウント
from google.colab import drive
drive.mount('/content/drive')

## Google Drive から zip をコピー
!cp /content/drive/My\ Drive/my_horse2zebra.zip ./

## URL をダウンロードするところを extract に変更
## https://www.tensorflow.org/datasets/api_docs/python/tfds/download/DownloadManager#extract
!git clone --depth 1 https://github.com/tensorflow/datasets
!perl -i.org -pe 's{data_dirs = .*}{data_dirs = dl_manager.extract("/content/my_horse2zebra.zip")};' datasets/tensorflow_datasets/image_classification/cycle_gan.py
!diff datasets/tensorflow_datasets/image_classification/cycle_gan.py{.org,}
!cd datasets/; pip install -e .

## import しておく
import sys
sys.path.append("/content/datasets/")
import tensorflow_datasets as tfds
print(tfds.__version__)

2. 学習

「Checkpoints」では checkpoint_path に前回のチェックポイントがあれば自動的に読み込んでくれるので (デフォルトだと ./checkpoints/train = /content/checkpoints/train)、Google Drive 上のディレクトリのシムリンクを作っておきます。

!mkdir -p /content/drive/My\ Drive/cyclegan_checkpoints/
!ln -s /content/drive/My\ Drive/cyclegan_checkpoints/ /content/checkpoints

後は cyclegan.ipynb をただ上から実行すれば学習・出力されます。

前提

ビデオツール部分の実装などは行っていません。

ビデオツールの背景機能ってどうやってるんだ？

ZoomやGoogle Meetなどで実装されている背景機能は自分で実現できないだろうかと思って。
React.jsとTensorflow.jsを使って実装してみることにした。

完成物

逆に自分を白く塗りつぶした。

クライアント完結で機械学習を使える！

機械学習のアプリを作ろうと思うと、Pythonでテンプレートを使って書くか、PythonでREST APIを書いて、React.jsなどでフロント側を作るという方法しか知らなかった。
個人的にはReact.jsを使い慣れているので、UIがある時にはReact.jsを使いたい。

そんな時、Tensorflow.jsを知った！
Tensorflow.jsを使うことで、Pythonを書くことなく、機械学習を使ったアプリが作れる！

すぐに使えるモデルがたくさん

Tensorflow.jsでは事前トレーニング済みのモデルが用意されている。
詳しくはこちら。

今回は人体セグメンテーション（BodyPix）のモデルを使用して背景機能を作っていく。

準備

最初に

$ npx create-react-app bodypix-web
$ cd bodypix-web

必要ライブラリのインストール

$ yarn add @tensorflow-models/body-pix @tensorflow/tfjs # 機械学習用
$ yarn add react-webcam # カメラ用
$ yarn add @material-ui/core # UI用

BodyPixを使って人体セグメンテーション

1. モデルのロード
2. セグメンテーションの生成
3. セグメンテーションの描画

モデルのロード

@tensorflow-models/body-pixのload()メソッドがあり、呼び出すことでロードができる。

import React, { useState, useEffect } from "react";
import "@tensorflow/tfjs";
import * as bodyPix from "@tensorflow-models/body-pix";

const App = () => {
    const [model, setModel] = useState();

    useEffect(() => {
        bodyPix.load().then((net) => {
            setModel(net);
        })
    }, []);
    return ();
}

また、load()の引数として以下のようなオプションを与えることができる。

• architecture: セグメンテーションのモデルを指定できる。MobileNetV1やResNet50が指定できる（デフォルトはMobileNetV1）。
• outputStride: セグメンテーションのモデルの出力ストライドを指定できる。値が小さいほど出力される解像度が高くなり、モデルの精度が向上するが速度は遅くなる。
• multiplier: MobileNetV1のみで使用できる畳み込みの深さを指定できるパラメータ。値が大きいほど精度が向上するが速度は遅くなる。
• quantByes: 重みの量子化に使用するバイト数を指定できる。何ができるかあまりわかっていない。

セグメンテーションの生成

ロードしたモデルのsegmentPerson()メソッドを使うことでセグメンテーションを生成することができる。

import React, { useState, useEffect } from "react";
import { Button, Box } from "@material-ui/core";

import "@tensorflow/tfjs";
import * as bodyPix from "@tensorflow-models/body-pix";

import WebCam from "react-webcam";

const App = () => {
    const [model, setModel] = useState();

    useEffect(() => {
        bodyPix.load().then((net) => {
            setModel(net);
        })
    }, []);


    // ここに注目
    const estimate = useCallback(() => {
        const webcam = document.getElementById("webcam");
        model.segmentPerson(webcam).then((segmentation) => {
            console.log(segmentation);
        }, [model])
    });

    return (
        <>
            <Box>
                <Button onClick={estimate}>推定</Button>
            </Box>
            <Box>
                <WebCam id="webcam" width={640} height={480} />
            </Box>
        </>
    );
}

segmentPerson()の引数としてはImageData、HTMLImageElement、HTMLCanvasElement、HTMLVideoElementが指定できる。今回はHTMLVideoElementを指定している。

セグメンテーションの描画

@tensorflow-models/body-pixのtoMask()メソッドでセグメンテーションに対応したピクセル値を生成し、drawMask()メソッドでそのマスクを画像やキャンバスに描画する。

import React, { useState, useEffect } from "react";
import { Button, Box } from "@material-ui/core";

import "@tensorflow/tfjs";
import * as bodyPix from "@tensorflow-models/body-pix";

import WebCam from "react-webcam";

const App = () => {
    const [model, setModel] = useState();

    useEffect(() => {
        bodyPix.load().then((net) => {
            setModel(net);
        })
    }, []);


    const estimate = useCallback(() => {
        const webcam = document.getElementById("webcam");
        model.segmentPerson(webcam).then((segmentation) => {
            showResult(segmentation);    
        }, [model, showResult])
    });

    // ここに注目
    const showResult = useCallback((seg) => {
        const foregroundColor = { r: 0, g: 0, b: 0, a: 0 };
        const backgroundColor = { r: 127, g: 127, b: 127, a: 255 }; 

        const mask = bodyPix.toMask(seg, foregroundColor, backgroundColor);

        const webcam = document.getElementById("webcam");
        const canvas = document.getElementById("canvas");
        const opacity = 0.7;

        bodyPix.drawMask(canvas, webcam, mask, opacity, 0, false);
    }, );

    return (
        <>
            <Box>
                <Button onClick={estimate}>推定</Button>
            </Box>
            <Box>
                <WebCam id="webcam" width={640} height={480} />
                <canvas id="canvas" width={640} height={480} />
            </Box>
        </>
    );
}

描画するためにまずtoMask()メソッドを使用しマスクを作成する。toMask()メソッドでマスクを作るためには前景と背景の色を決める必要がある。ここではrgba指定で作成し、toMask()メソッドにセグメンテーション結果とともに渡している。

次に、drawMask()メソッドを使用してキャンバスに描画する。drawMask()メソッドは以下の引数が必要となる。

• canvas: キャンバスの要素。
• image: 適用する元画像。VideoElementも指定できるため、webcamを指定している。
• maskImage: toMask()などで生成されたマスクデータ。
• maskOpacity: 画面上にマスクを描画する際に不透明度。デフォルトは0.7となっている。
• maskBlurAmount: マスクをぼかすピクセルすう。0~20が指定できる。デフォルトは0となっている。
• flipHorizontal: 結果を反転させるかどうかを指定できる。デフォルトはfalseとなっている。

他にやったこと

• リアルタイム描画
• 背景塗りつぶし、人物塗りつぶし、塗りつぶしなしの切り替え

コードはこちら
https://github.com/makky0620/bodypix-web

最後に

サンプルコードとかではReact.jsを使った例がなかったので、ちゃんとかけているか分からないですが、結構簡単に実装できて嬉しかった。
他のモデルも触ってみよう！