kivy+PyTorchでお絵かきクイズを作る

あらまし

この記事はドコモアドベントカレンダー2日目の記事です．

お絵かきクイズってご存知ですか？出題者が絵を描いて，他の人が何を描いているか当てるクイズです．アプリもたくさん出ています．

ユーザが描いたものを画像認識して当ててもらえば一人でも遊べそうなので，去年に引き続き機械学習ライブラリのPyTorchとアプリを作成できるライブラリのkivyを組み合わせて作ってみます．

kivyとは？

PythonのGUIライブラリの中で最も勢いがあるし使いやすいです（主観）．

日本語の情報がないので敷居が高いと言われていた時代もありましたが，今はめちゃくちゃ情報があります．有志のおかげで日本語ドキュメントもあります．はじめての人向けチュートリアルとしてはこの記事がおすすめです．

• GOOD
  • MITライセンスで商用利用可能（よく比較されるPyQtは商用は有償）
  • マルチプラットフォーム対応（PC向け/スマホ向けの両方が作れる）
• BAD
  • UI部分は独自言語（kivy言語）で記載する必要がある
    • MarkdownとかXMLくらい直感的に読み書き出来ます

検証環境

• Windows10 64bit
  • OSによってkivyのインストール方法が違います
• Python 3.7.3
• kivy 1.11.1
• torch 1.0.1

レイアウトの策定

下記の機能があればお絵かきクイズが実現できそうです．

• 描くべきお題の表示
• お題を変更するボタン
• キャンバス部分
• キャンバスの消去ボタン
  • 全消し
  • 直前の線を消す（Undo）
• 画像認識の結果の表示

kviewerを使うと本体のpythonファイルがなくてもUIの確認が出来るので先にレイアウトだけ作ってみます．

ui_test.kv

BoxLayout:
    orientation: 'vertical' # 4つの子オブジェクトを縦に並べる
    BoxLayout:
        size_hint_y: 0.1 # 縦幅の10%を占める
        Label: # お題を表示
            size_hint_x: 0.7 # 横幅の70%を占める
            text: 'Draw "XXX"'
        Button: # お題を変更するボタン
            size_hint_x: 0.3 # 横幅の30%を占める
            text: 'Change'
    Label: # キャンバス部分（仮）
        text: 'CANVAS'
        size_hint_y: 0.7 # 縦幅の70%を占める
    BoxLayout:
        size_hint_y: 0.1 # 縦幅の10%を占める
        Button: # Undoボタン
            text: 'Undo'
        Button: # 全消しボタン
            text: 'Clear'
    Label: # AIが絵を当てた結果を表示
        size_hint_y: 0.1 # 縦幅の10%を占める
        text: 'I guess it is "YYY"'

表示すると下記のような画面が出てきました．レイアウトはこれで決定です．

python -m kivy.tools.kviewer ui_test.kv

認識モデルの用意

ユーザが描いた絵を当てるモデルが必要ですが，今回は下記の学習済みモデルを使います（MIT License）．軽い代わりに20カテゴリしかないので，元のデータセットにある全345カテゴリを認識させたい人は自前で学習してみてください．

• https://github.com/nhviet1009/QuickDraw
  • GoogleのQuick Draw Datasetの一部で学習されたモデルです．

今回使うモデルは時系列情報（線を描いた順序）は用いていませんが，元のデータセットには時系列情報もあるのでこれを考慮して学習すると描いている途中での認識精度が向上するはずです．
フォルダ構造は下記のようにします．上記のGitHubのリポジトリから「src」「trained_models」の2フォルダを使います．

├─paintquiz.py # これから実装するpythonファイル
├─paintquiz.kv # これから実装するkivyファイル
├─src # 上記リポジトリ内
└─trained_models # 上記リポジトリ内

実装

pythonで必要な機能を実装していき，それに合わせてkivy側も少し変更します．pythonが約120行，kivyが約60行，計約180行で実装できました．
1秒ごとにキャンバスを認識して暫定の認識結果を表示するようにしています．実装した機能は下記です．

• レイアウトで定めた各ボタンの機能
  • Changeボタン（お題の変更）
  • Undoボタン（直前の線を消す）
  • Clearボタン（キャンバス全消し）
• キャンバスへの描画機能
• キャンバスに描かれたものをリアルタイムで認識する機能

Pythonコード

paintquiz.py

import random

from kivy.config import Config
# 起動時のウィンドウサイズ
Config.set('graphics', 'width', '600')
Config.set('graphics', 'height', '700')

from kivy.app import App
from kivy.clock import Clock
from kivy.properties import StringProperty
from kivy.uix.widget import Widget
from kivy.uix.boxlayout import BoxLayout
from kivy.uix.label import Label
from kivy.graphics import Line, Color

import numpy as np
import cv2
import torch
import torch.nn as nn

# 認識できるクラス一覧
classes = ["apple", "book", "bowtie", "candle", "cloud", "cup", "door", "envelope", "eyeglasses", "guitar", "hammer",
           "hat", "ice cream", "leaf", "scissors", "star", "t-shirt", "pants", "lightning", "tree"]

class Net():
	# 認識モデルを読み込む
    def __init__(self):
        self.model = torch.load("./trained_models/whole_model_quickdraw", map_location=lambda storage, loc: storage) # 学習済みの重みファイルを読み込む
        self.model.eval() # 今回は学習は行わない
        self.sm = nn.Softmax(dim=1) # どの結果も信頼性が低い場合にスコアで足切りしたいのでsoftmaxで正規化する

	# 画像ファイル名を入力して認識結果を返す
    def predict(self, fn, th=.5):
        image = cv2.imread(fn, cv2.IMREAD_UNCHANGED)[:,:,-1] # alpha channelを取得して2値画像へ
        image = cv2.resize(image, (28, 28))
        image = np.array(image, dtype=np.float32)[None, None, :, :]
        image = torch.from_numpy(image)
        pred = self.model(image)
        pred = self.sm(pred)
        return torch.max(pred[0]), torch.argmax(pred[0]) # 認識スコアと認識クラスを返す

class Paint(Widget):
    pred_word = StringProperty() # 認識結果の単語
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.line_width = 10 # 線の太さ
        self.lines = [] # undo用に線を格納するリスト
        self.in_drawing = False # 描画中か否かを判定
        self.canvas.add(Color(0,0,0))
        self.model = Net()
        Clock.schedule_interval(self.predict, 1.0)

    def calc_pos(self, bbox):
        xmin = min(bbox[0], bbox[2])
        ymin = min(bbox[1], bbox[3])
        xmax = max(bbox[0], bbox[2])
        ymax = max(bbox[1], bbox[3])
        return xmin,ymin,xmax,ymax

    # クリック中（描画中）の動作
    def on_touch_move(self, touch):
        if self.in_drawing == False:
            if self.pos[0]<touch.x<self.pos[0]+self.size[0] and self.pos[1]<touch.y<self.pos[1]+self.size[1]:
                self.in_drawing = True
                with self.canvas:
                    touch.ud['line'] = Line(points=(touch.x, touch.y), width=self.line_width)
        elif touch.ud:
            if self.pos[0]<touch.x<self.pos[0]+self.size[0] and self.pos[1]<touch.y<self.pos[1]+self.size[1]:
                touch.ud['line'].points += [touch.x, touch.y]

    # クリック終了時の動作
    def on_touch_up(self, touch):
        if self.in_drawing:
            self.lines.append(touch.ud['line'])
            self.in_drawing = False

    # 直前の線を消去
    def undo(self):
        if len(self.lines)>0:
            line = self.lines.pop(-1)
            self.canvas.remove(line)

    # キャンバスを全消しする
    def clear_canvas(self):
        for line in self.lines:
            self.canvas.remove(line)
        self.lines = []

    # dt秒ごとに画像を認識する
    def predict(self, dt):
        self.export_to_png('image.png')
        with torch.no_grad():
            score, label = self.model.predict('./image.png')
            # 認識スコアが一定以下の場合にはわからないと表示
            if score < 0.5:
                self.pred_word = "CPU: I have no idea"
            else:
                self.pred_word = 'CPU: I guess it is "{}"'.format(classes[label].upper())

class PaintQuiz(BoxLayout):
    word = StringProperty('Draw "{}"'.format(random.choice(classes).upper())) # お題の単語
    def __init__(self, **kwargs):
        super(PaintQuiz, self).__init__(**kwargs)
        pass

    def reset(self):
        self.word = 'Draw "{}"'.format(random.choice(classes).upper())

class PaintQuizApp(App):
    def __init__(self, **kwargs):
        super(PaintQuizApp, self).__init__(**kwargs)
        self.title = 'PAINT QUIZ'

    def build(self):
        return PaintQuiz()

if __name__ == '__main__':
    app = PaintQuizApp()
    app.run()

Kivyコード

paintquiz.kv

<PaintQuiz>:
    canvas:
        Color:
            rgb: .9,.9,.9
        Rectangle:
            pos: self.pos
            size: self.size
    BoxLayout:
        size: root.size
        orientation: 'vertical'
        
        BoxLayout:
            size_hint_y: 0.1
            orientation: 'horizontal'
            Label:
                canvas.before:
                    Color:
                        rgb: 1.,.3,.3
                    Rectangle:
                        pos: self.pos
                        size: self.size
                size_hint_x: 0.7
                text: root.word
                font_size: 18
            Button:
                id: button_reset
                size_hint_x: 0.3
                text: 'Change'
                on_release: root.reset()
        Paint:
            size_hint_y: 0.7
            id: paint_area
            allow_stretch: True
        BoxLayout:
            size_hint_y: 0.1
            Button:
                id: button_undo
                text: 'Undo'
                on_release: paint_area.undo()
            Button:
                id: button_clear
                text: 'Clear'
                on_release: paint_area.clear_canvas()
        Label:
            canvas.before:
                Color:
                    rgb: 1.,.3,.3
                Rectangle:
                    pos: self.pos
                    size: self.size
            size_hint_y: 0.1
            text: paint_area.pred_word
            font_size: 18

実行結果

左上に現在のお題，下に認識結果（何を描いていると思っているか）が表示されます．

リンゴやハサミを描いてみましたが正しく認識してくれています．リンゴはヘタだけ描いた時点だとハンマーや葉っぱに見えますが，確かにハンマーや葉っぱとして認識されています．むしろリンゴのヘタ＝葉っぱなので正解ですね．

おわりに

データ分析界隈だと打ち合わせのときにパワポじゃなくてJupyter Notebookで結果を見せる人がいますが，そのくらいの感覚でアプリが作れるし興味を惹きやすいのでおすすめです．

本当は絵しりとりアプリ（AIと交互に絵でしりとりをする）を作りたかったのですが，kivy経験者限定！な記事になりそうだったので断念しました．誰か作って遊ばせてください．