CNNを活用しTwitterの画像欄からソシャゲと写真を避けてイラストの抽出を行う

はじめに

twitterでめちゃくちゃ好みのイラストが流れて来た時に、そのイラストを描いてる人のほかの作品が見たくなって画像欄を見に行ってみるとソシャゲのスクショや飯テロ画像が多く、なかなか他のイラストを見つけることができない…ということが多々ありました。
そこで、ちょうど機械学習を勉強したいという欲もあったので、CNNを活用することでイラストとソシャゲのスクショ、写真を識別し、イラストのみを抽出することができないかと思って試してみました。

手順

上記のことを達成するために、今回は以下の順番で抽出方法を考えていきました。

1. 識別方法を考える
2. 画像の収集
3. 学習
4. 学習したモデルの再評価
5. 考察

識別方法を考える

イラストの抽出するにはどうすればいいかな・・・と考えながらtwitterに流れてくる画像を眺めていると、イラストのほかにソシャゲなどのスクショや写真の画像が多かったように感じました。そこで、イラスト、スクショ、写真の3クラスに分類する識別機を作成することでうまくイラストの画像だけ抜きだせるのではと思い、CNNを活用した3クラスへの分類を試みることにしました。

画像の収集

今回、3クラスに分類を試みることを決めたので、それぞれに対応する教師画像を集めることにしました。
集め方はTwitterのAPIを使って画像ツイートを収集し、その後手動でクラス分けを行っていきました。
あるユーザーが投稿した画像を取れるだけとってくるものとなっています。
収集対象としては、イラストを描いている人からランダムで10人を選び、そのユーザーが投稿した画像を収集しました。

次にクラス分けなのですが、主観が入ってクラス分けの基準がぶれてしまうことをできるだけ避けるために、あらかじめそれぞれのクラスに振り分ける基準を作り、それにしたがってクラス分けを行っていくことにしました。
基準は以下のように決めました。

• イラスト
  画像ツイートの投稿主が描いたもの
• スクショ
  ソシャゲのスクリーンショット。ゲーム途中の画面、ガチャの結果などすべてここに含める。
• 写真
  カメラで現実の光景を撮影したもの

基準を決めてクラス分けを行っていたのですが、漫画の一コマを切り取ったものといった想定していない画像も出てきてしまいました。今回はそれらは教師データに含めないという処理をしたのですが、次はちゃんと分けられるようにしたいです。
収集した画像は各クラスで400枚ほど集めました。

学習

教師データが集まったので学習を行っていきイラストを識別できるようなモデルを生成していくことにしました。
今回集まった画像は各クラスで400枚であり、合計1200枚程度であり、CNNのモデルを生成するにはかなり枚数が少ないと感じます。
そこで、今回はfine-tuningという少量のデータでも十分な識別性能が得られる手法を採用しました。
fine-tuningとは大量のデータを用いて学習した別タスクのモデルを流用して自分が用意したクラスに識別できるように再学習する技術です。すでに十分な表現力を持ったモデルを使うので十分な識別性能が得られる手法とのことです。
pytorchのチュートリアル(https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html )にfine-tuningの実施方法を記述してあるものがあったのでそれを使わせていただきました。ここでは、パラメータチューニングは行っていないです。
一応おのおののパラメータが何を意味しているのかは読みましたが、ちゃんと説明できるほど理解はしていないという状態です…ちゃんと学習していきたい…
チュートリアルに詳しい説明が書いてあったので、ここでは細かい説明は省略したいと思います。
このチュートリアルでは25回のエポックのうち、精度が一番よかったモデルのパラメータを採用するようになっています。

training.py

# License: BSD
# Author: Sasank Chilamkurthy

from __future__ import print_function, division

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
import numpy as np
import torchvision
from torchvision import datasets, models, transforms

data_transforms = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'val': transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}
data_dir = 'D:\\data\\twitter_train\\'

image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),
                                          data_transforms[x])
                  for x in ['train', 'val']}
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=4,
                                             shuffle=True, num_workers=4)
              for x in ['train', 'val']}
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
class_names = image_datasets['train'].classes

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model_ft = models.vgg16(pretrained=True)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs, 3)

model_ft = model_ft.to(device)

criterion = nn.CrossEntropyLoss()

optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)

def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
    since = time.time()

    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0

    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 10)

        # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
            if phase == 'train':
                scheduler.step()
                model.train()  # Set model to training mode
            else:
                model.eval()   # Set model to evaluate mode

            running_loss = 0.0
            running_corrects = 0

            # Iterate over data.
            for inputs, labels in dataloaders[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)

                # zero the parameter gradients
                optimizer.zero_grad()

                # forward
                # track history if only in train
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)

                    # backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                # statistics
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)

            epoch_loss = running_loss / dataset_sizes[phase]
            epoch_acc = running_corrects.double() / dataset_sizes[phase]

            print('{} Loss: {:.4f} Acc: {:.4f}'.format(
                phase, epoch_loss, epoch_acc))

            # deep copy the model
            if phase == 'val' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())

        print()

    time_elapsed = time.time() - since
    print('Training complete in {:.0f}m {:.0f}s'.format(
        time_elapsed // 60, time_elapsed % 60))
    print('Best val Acc: {:4f}'.format(best_acc))

    # load best model weights
    model.load_state_dict(best_model_wts)
    return model

model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler,num_epochs=25)

data_dir = 'D:\data\twitter_train\'

ここには教師データの格納したディレクトリを指定しています。
ディレクトリ内は各クラスごとのディレクトリを作成する必要があります。今回は以下のようなディレクトリ構造にしました。
twitter_train/
　├ illust/ <イラストを格納>
　├ picture/ <写真を格納>
　├ screenshot/ <スクショを格納>

実行した結果、識別精度は94.2%が得られました。

学習したモデルの評価

学習した精度が94.2%となったのですが、ほんとにいろんなイラストの抽出ができるか気になりました。
そこで、学習に使ったものとは別に新たに各クラスごとの画像を収集し、クラス分けが正しく行えているのかを評価しました。
特にイラストは描いている人によって画像の特徴が異なっており、正しくイラストが抽出できない可能性が高そうだな…と思ったので、できるだけ違う絵柄の人を選んで評価に使いました。（ここは主観が入った選定になってしまいました）
今回の評価では、スクショや写真がうまく分けられないことより、イラストの画像を抽出できることが重要と考え、精度のほかにイラストクラスについての再現率という指標も採用して評価することにしました。
今回の評価に使ったデータセットとしては各クラスに100枚の画像を用意して識別を行っていきました。
結果としては、識別率：75%, 再現率は85%という結果になりました。

考察

やはり描いている人によって特徴があるせいか、評価する画像を変えると20%近くも精度が落ちる結果となりました。
ただ再現率は85%あり、実際にイラスト抽出システムに利用できそうだなと感じます。ただ個人的にはイラストを見逃してしまう可能性をできるだけ下げたいのでもう少し高くしたいところです…
そこで、誤識別を行ってしまった画像はどのような画像が多かったのかを詳細に見て改良できるポイントがないかを見てみました。
まず、イラストなのに、写真やスクショと判断されてしまった画像を見てみました。
以下に一例を示します。ただ、そのまま載せるとまずいかなと思いぼかしたものを載せています。

見てみると、イラストの背景がとてもきれいに描かれているおり、写真のように見えます。このような画像はイラストではなく、写真と判断されてしまうようなモデルになっているようです。
学習データに使っていたのは漫画のようなイラストが多く、このような結果になったのかなと考えています。
次に学習データを作る際にはイラストの種類を定義づけて様々なデータを含むようにしてみたいと思います。
イラストの種類を分別できるようにするのもかなり面白そうです！

また、イラストではないのにイラストと判断されてしまっているものにどのようなものがあるのかも見てみました。

紙などに描かれたものを写真で読み取ったものはイラストと判定していました。今回はカメラで撮影したものをすべて写真と定義したためこのような結果になっているのですが、もともとの目的はイラストを画像群から抜きだすことであり定義次第では目的の画像であるといえるのでモデルとしては問題ないのかなと思います。
amazonの商品ページがイラストと判断されてしまっているのはよくわからない結果でした。。何をもってイラストと判定する判別機になっているのか気になるところです。。
ただ、当初懸念していたソシャゲのスクショがイラストと判断されるケースは今回は見当たりませんでした。
学習に入れてなかったソシャゲのスクショやガチャ結果などのキャラクターだけを映したスクショも含まれていたのですが、イラストと判断されなかったのは意外でした。
この結果から考えると、ソシャゲは似ている構図や画面になっており、それらを識別可能な特徴量が抽出できているのではないか？と考えています。

まとめると、イラストの識別は学習に使ったイラストに依存する可能性が高そうということがわかりました。イラストを取り逃すことなく抽出するためには一人から抽出したイラストを使うのではなく、できるだけ多くの人から画像を抽出して使用することが重要になりそうです。
ソシャゲのスクショや写真といった目的のイラストではない画像が混入することはここまでの評価を見る限り少ないと考えられそうでした。

まとめ

CNNを使って、イラストの画像を抽出するということを行ってみました。
結果的にはイラスト欄を眺めた時にソシャゲの画像や飯テロ画像がほとんど現れずに画像を閲覧することができました！
使っている技術としては特に新しいものを使っていなかったのですが、自分が目指していたことはある程度達成できたかなと感じています。CNNすごい。
今回はチュートリアルの内容を少しだけ変えて実行していっただけだったのですが、自分が何をしたいのかを考えながら実行していくのは楽しかったので、今後も機械学習系の学習は続けていきたいなと思います。
今度はもう少し難しめの内容にもチャレンジしてみたいです！