機械学習に用いるデータセット作成のための画像収集Pythonスクリプト

背景

ある問題に対して機械学習を用いて高精度の結果を出力させるためには、データセットの作成と学習器の設計をどのように工夫するかが重要となる。この内、データセット作成においてデータ数をどれだけ用意するかということは、学習を成功させる上で特に重要となってくる。
学習に必要なデータを収集するためにまず行うことが、既に世の先人達がWeb上に溢れているデータや自前で作成したデータを整理して構築してくれたデータベースを探すことである。しかし、そんな都合の良いデータベースが毎回あるとは限らないため、自分で収集してくる必要があることもある。
そこで今回は、勉強のために簡単な例として人の画像を収集した上で、顔のみをCropして保存するスクリプトをPythonで作成した。

実行環境

• Mac OS X 10.10.5 (Yosemite)
• Python 2.7.13_0
• Open CV 3.2.0_0 (+contrib+java+python27+qt4+vtk)

目的

今回作成するスクリプトは、クエリを与えたら関連する画像をWeb上で検索して収集し、OpenCVを用いて顔認識を行った上でCropして保存するというものを想定している。また、用いた検索エンジンはBingの画像検索である。

スクリプト

本スクリプトのソースコードは以下のリンク先に置いてある。
https://github.com/tokkuman/FaceImageCollection
以下からは、本スクリプトにおける各関数について説明する。

Import

今回のスクリプト作成におけるインポートしたモジュール群は以下のようになっている。

import sys
import os
import commands as cmd
import cv2
import time
import copy
from argparse import ArgumentParser

getHtml

クエリを投げて検索したページのhtmlを取得する関数。
cmd.getstatusoutputはタプルで(status, output)を返すため、今回はwget -Oで引っこ抜いてきたhtmlのみをreturnしている。

def getHtml(query):
    return cmd.getstatusoutput("wget -O - https://www.bing.com/images/search?q=" + query)[1]

extractImageURL

htmlと画像のformatを受け取り、htmlから指定したformatのリンクのみを抽出してくる関数。

def extractImageURL(html, suffix):
    url = []
    snum, lnum = 0, 0
    text = html.split('\n')
    for sen in text:
        if sen.find('<div class="item">') >= 0:
            element = sen.split('<div class="item">')
            for num in range(len(element)):
                for suf in suffix:
                    snum = element[num].find("href") + 6
                    lnum = element[num].find(suf) + len(suf)
                    if lnum > 0:
                        url.append(element[num][snum:lnum])
                        break
    return url

saveImg

extractImageURLで抽出してきたリンクからお目当の画像をいったんローカルに保存する関数。
作成したディレクトリ(opbase)内にさらにOriginalというディレクトリを作成して、その中に画像を保存する。

def saveImg(opbase, url):
    dir = opbase + '/Original/'
    if not (os.path.exists(dir)):
        os.mkdir(dir)
    for u in url:
        try:
            os.system('wget -P ' + dir + ' '  + u)
        except:
            continue

cropFace

保存した画像の中から、顔のみを抽出して切り抜き保存する関数。
顔認識は、OpenCVにおけるHaar分類器の学習済みモデル(haarcascade)をロードして用いている。本スクリプトでは正面からの顔のみを抽出することを想定して、4種類のmethodを用いることを可能としている。モデルの精度に関しては、以下のリンク先 (http://stackoverflow.com/questions/4440283/how-to-choose-the-cascade-file-for-face-detection) などを参考にした。
Crop後の画像はCropディレクトリをopbase内に作成して、その中に保存する。

def cropFace(opbase, path, imsize, method):
    dir = opbase + '/Crop/'
    if not (os.path.exists(dir)):
        os.mkdir(dir)
    for p in path:
        img = cv2.imread(opbase + '/Original/' + p)
        gImg = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        if method == 1:
            face_cascade = cv2.CascadeClassifier('haarcascade/haarcascade_frontalface_default.xml')
        elif method == 2:
            face_cascade = cv2.CascadeClassifier('haarcascade/haarcascade_frontalface_alt.xml')
        elif method == 3:
            face_cascade = cv2.CascadeClassifier('haarcascade/haarcascade_frontalface_alt2.xml')
        else:
            face_cascade = cv2.CascadeClassifier('haarcascade/haarcascade_frontalface_alt_tree.xml')
        faces = face_cascade.detectMultiScale(gImg, 1.3, 5)
        for num in range(len(faces)):
            cropImg = copy.deepcopy(img[faces[num][1]:faces[num][1]+faces[num][3], faces[num][0]:faces[num][0]+faces[num][2]])
            resizeImg = cv2.resize(cropImg, (imsize, imsize))
            filename = dir + p[:-4] + '_' + str(num + 1) + '.tif'
            cv2.imwrite(filename, resizeImg)

main

本スクリプトのmain関数。
parserで以下のオプションを指定して実行する。

• query : 検索するクエリを指定。クエリは複数指定可能。
• suffix : 保存する画像の形式を指定。同じく複数指定可能。
• imsize : Cropした顔画像のサイズを指定。
• method : Haar分類器におけるモデルの種類を指定。

また、指定したクエリ名で本スクリプトの出力ディレクトリが作成される。

if __name__ == "__main__":
    ap = ArgumentParser(description='ImageCollenction.py')
    ap.add_argument('--query', '-q', nargs='*', default='hoge', help='Specify Query of Image Collection ')
    ap.add_argument('--suffix', '-s', nargs='*', default='jpg', help='Specify Image Suffix')
    ap.add_argument('--imsize', '-i', type=int, default=100, help='Specify Image Size of Crop Face Image')
    ap.add_argument('--method', '-m', type=int, default=1, help='Specify Method Flag (1 : Haarcascades Frontalface Default, 2 : Haarcascades Frontalface Alt1, 3 : Haarcascades Frontalface Alt2, Without : Haarcascades Frontalface Alt Tree)')
    args = ap.parse_args()

    t = time.ctime().split(' ')
    if t.count('') == 1:
        t.pop(t.index(''))
    # Path Separator
    psep = '/'
    for q in args.query:
        opbase = q
        # Delite File Sepaeator   
        if (opbase[len(opbase) - 1] == psep):
            opbase = opbase[:len(opbase) - 1]
        # Add Current Directory (Exept for Absolute Path)
        if not (opbase[0] == psep):
            if (opbase.find('./') == -1):
                opbase = './' + opbase
        # Create Opbase
        opbase = opbase + '_' + t[1] + t[2] + t[0] + '_' + t[4] + '_' + t[3].split(':')[0] + t[3].split(':')[1] + t[3].split(':')[2]
        if not (os.path.exists(opbase)):
            os.mkdir(opbase)
            print 'Output Directory not exist! Create...'
        print 'Output Directory:', opbase

        html = getHtml(q)
        url = extractImageURL(html, args.suffix)
        saveImg(opbase, url)
        cropFace(opbase, os.listdir(opbase + '/Original'), args.imsize, args.method)

結果

どのくらいノイズが混じるか実験するため、今回は個人的な趣味で"ガッキー"と"ベッキー"というクエリを投げて得られた結果を示す。

さすがは天下のガッキーと言わざるを得ないというべきか、一部化け物ノイズも含まれているが概ねガッキーだった。一方ベッキーは、ベッキー以外のベッキーと思われる人もとってきてしまっている。一般性の高いクエリほどノイズが多く含まれてしまうことは仕様上致し方ないが、改善の余地があるといえる。また、一番改善すべき点は収集できた画像数であり、1クエリに対する画像数が圧倒的に少ないことから、まだまだ全体的に課題が残っているといえる。

考察

世の中的に、画像収集はGoogle Custom Search APIやBing Search APIなどを用いることが一般的であり、精度も収集数も圧倒的に高い。今回は、あえてそれらのAPIを用いずにどこまで試せるかというチャレンジでもあったが、APIを用いた方法も試してみたい。
また、htmlを解析する方法も探り探りだったため、抽出方法にも問題があったと考えられる。便利なモジュールとしてよくみかけるBeautifulSoupなどを用いてもいいかもしれない。
また、ガッキーは恋ダンスなどもあり大躍進している(もともと天使だから飛んでいた訳だが)一方、ベッキーは崖っぷちからの再スタートをきったという現状が、今回の結果に反映したとも考えられる。是非あいのりのMCで頑張っていただきたい。