【実例付き】PyTorchを使って作業を自動化してみよう

今回はPyTorchを使って簡単な作業を自動化するやり方を、実際のコードを使用しながら解説をしていきます。

お品書き

前提
自動化について
まとめ

前提

前提として、下記ソフトウェア・ライブラリがインストールされていることを確認してください。
また、ここではWindows上で動作している前提で解説をしていきます。

Python 3.9.7以上
PyTorch 1.10.1+cu102以上
shutil
glob

自動化について

今回自動化していくのはこのような作業になります。

ある日Aくんのお父さんは家族で花見に行きました。
Aくんのお父さんはカメラを持っていき、大量に写真を撮りました。
後からアルバムにしようと考えましたが、Aくんのお父さんは人と桜のアルバムを分けようと考えました。
ですが、数百枚あるデータを一からすべて人物が含まれるかどうかを手作業で確認するのは大変です。
そこで、Aくんのお父さんはなんとかして自動化できないかと考えました。

Aくんのお父さんは以下の様な形(花見用フォルダに画像を保存)で配置しました。
注意: ここではダミー用の画像(YOLOv5のサンプル画像)を使用します。

PyTorchにモデルを読み込む

まず、今回は人物が含まれるかどうかを判別に使うため、PyTorch上で人物のモデルを読み込みます。
ここで使用するモデルデータは、すでに学習されているYOLOv5のモデルを使用します。
PyTorchではtorch.hub.loadメソッドでモデルデータを読み込むことができます。

import torch
import glob
# モデルを読み込む
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

実際に実行するとYOLOv5のモデルがダウンロードされ、読み込まれた旨のログが出力されます。

YOLOv5  2022-1-4 torch 1.10.1+cu102 CUDA:0 (NVIDIA GeForce RTX 2080 SUPER, 8192MiB)

Fusing layers... 
Model Summary: 213 layers, 7225885 parameters, 0 gradients
Adding AutoShape...

画像を読み込む

では、実際にフォルダに入っている画像をすべて読み込み、人物が含まれている画像だけを抽出してみましょう。
globというパッケージを使用することでファイルパスを読み込み、for文で繰り返し処理ができます。
下記コードを入力してみましょう。

import torch
import glob
#モデルを読み込む
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

# 画像ファイルをすべて読み込む
for image in glob.glob('./花見用/*.jpg'):
    # 画像を読み込む
    print(image)

実際に実行するとファイルのパスが出力されます。

~~~ 省略 ~~~
Adding AutoShape... 
./花見用\img (1).jpg
./花見用\img (10).jpg
./花見用\img (11).jpg
./花見用\img (12).jpg
./花見用\img (13).jpg
./花見用\img (14).jpg
~~~ 省略 ~~~

人物を判別する

では実際に読み込んだ画像をPyTorchのモデルで人物の判別をしてみましょう。
globで取得したファイルパスから画像を読み取るにはPyTorchのモデルデータの変数に画像を引数として渡してあげる必要があります。
モデルデータから人物だけをとるにはモデルデータから物体が検出された四角(Bounding Box)を取得し、その中からクラスデータと照らし合わせる必要があります。
Bounding Boxはモデルデータ.xyxy[]で取得することができ、この配列の中身は[[左上のx軸, y軸, 右下のx軸, y軸, 一致率, クラス番号], cudaデバイス]という形になっています。
下記コードで人物を検出することが可能です。

# 画像ファイルをすべて読み込む
for image in glob.glob('./花見用/*.jpg'):
    # 画像をモデルに渡す
    result = model(image)
    # 配列をループして結果を出力
    for i in result.xyxy:
        # 配列からBoundingBoxだけを取得
        for j in i:
            # クラスIDを取得し0(Person)以外は除外
            # [x1, y1, x2, y2, confidence, classid]
            if j[5] == 0:
                # ここで人物を検出したときの処理を書く

上記コードでは画像ファイルをすべて読み込んだ後に、モデルで画像を処理し、配列の処理をしています。

判別した物を処理する

判別したものを処理するために上段落の後に追記していきます。
今回は人物が含まれている画像をoutputディレクトリにコピーして保存することにします。
ファイルをコピーする際にはshutil.copy()関数を使用するため追記でimport shutilを記述します。

# 配列からBoundingBoxだけを取得
for j in i:
    # クラスIDを取得し0(Person)以外は除外
    # [x1, y1, x2, y2, confidence, classid]
    if j[5] == 0:
        # ここで人物を検出したときの処理を書く
        shutil.copy(image, 'output/')

実際に上記コードを実行すると、人物が含まれている画像をoutputディレクトリに全てコピーをすることができ、問題を解決することができます。