前回は学習データとして使う犬の顔部分のアノテーションの作成を行いました。
本記事では、いよいよモデルを学習し、目的である犬の顔に自動でモザイク処理をするAIの完成を目指します。
モデルの学習
DeepLearningモデルの学習を高速化するにはGPUは必要不可欠でしょう。そこで今回は無料でGPUを使うことができるGoogle Colaboratoryを使用してモデルの学習を行います。
Google Colaboratory自体の使い方については今回省略させていただきます。
Google ColaboratoryでGPUの設定をしているとして、次に進みます。
学習データのアップロード
前回まででプロジェクトフォルダは以下のようになっていると思います。
dog-face-mosaic-ai
|- env
|- models
| |- yolov5s.pt
|- data
| |- dog
| | |- flickr_dog_000043.jpg
| | :
| | :
| |
| |- annotation
|
|- main.py
|- requirements.txtdataフォルダを丸ごとGoogle Colaboratoryのディスク上にアップロードします。
(Google Driveにアップロードして、Google Colaboratoryとマウントするでも良いです。)
各種ライブラリのインポート
yolov5のリポジトリのcloneと使用するライブラリのインストールとインポートをします。
以下がコードです。
!git clone https://github.com/ultralytics/yolov5.git > /dev/null 2>&1
%cd /content/yolov5/
!pip install -qr requirements.txtimport os
import shutil
import yaml
import random
import glob
import tqdm学習データと検証データを分ける
学習させるデータと学習の効果を確認するための検証データを振り分けます。
まずdataフォルダまでのパスを変数に格納します。
こちらのパスは適宜データをアップロードした場所に変更してください。
data_path = "/content/drive/MyDrive/data"続いて振り分けた後に保存しておく各フォルダを作成します。
os.makedirs("/content/dataset/train/images", exist_ok=True)
os.makedirs("/content/dataset/train/labels", exist_ok=True)
os.makedirs("/content/dataset/val/images", exist_ok=True)
os.makedirs("/content/dataset/val/labels", exist_ok=True)そして以下コードを実行すると学習用と検証用が8:2の割合でランダムに振り分けられます。
label_paths = sorted(glob.glob(f'{data_path}/annotations/*'))
label_paths = [path for path in label_paths if not("classes.txt" in path)]
random.shuffle(label_paths)
for i, label_path in enumerate(tqdm.tqdm(label_paths)):
file_name = os.path.basename(label_path)[:-4]
img_path = glob.glob(f"{data_path}/**/{file_name}.jpg", recursive=True)[0]
if i < int(len(label_paths) * 0.8):
shutil.copy2(label_path, f"/content/dataset/train/labels/{file_name}.txt")
shutil.copy2(img_path, f"/content/dataset/train/images/{file_name}.jpg")
else:
shutil.copy2(label_path, f"/content/dataset/val/labels/{file_name}.txt")
shutil.copy2(img_path, f"/content/dataset/val/images/{file_name}.jpg")モデル学習実行
学習用の設定ファイルを作成します。
data = {
"train": "/content/dataset/train",
"val": "/content/dataset/val",
"nc": 1,
"names": ["dog"]
}
with open('/content/data.yaml', 'w') as f:
yaml.dump(data, f, default_flow_style=False)そして以下コマンドで学習を実行。
!python train.py --data "/content/data.yaml"今回は学習時の条件は特に設定していませんが、コマンドの引数に学習エポック数など設定できるので興味がありましたら調べてみてください。
学習が終わると /content/yolov5/runs/train/exp/weights/best.pt というファイルが作成されています。
これが学習結果の重みです。
こちらはダウンロードして控えておきましょう。
試しに犬の顔検出
まだGoogle Colabratoryで作業します。
学習したモデルが犬の顔を検出するか確認しましょう。
以下コマンドで推論を行います。
!python detect.py --source "/content/sample.jpg" --weights "/content/yolov5/runs/train/exp/weights/best.pt"/content/sample.jpg のところは適宜推論したい画像のパスに置き換えてください。
結果は以下コードで表示できます。
from PIL import Image
exp_name = "exp"
img = Image.open(f"/content/yolov5/runs/detect/{exp_name}/sample.jpg")
img
しっかりと犬の顔を検出するモデルを作成することができました。
犬の顔にモザイク処理をするAIの仕上げ
学習済みの重みさえゲットすれば、後は第2回の記事で作成したコードを使ってモザイク処理ができます。
プロジェクトフォルダに戻り、以下コマンドで動画と先ほど保存した学習済みの重みを選択することでモザイク処理ができます。
python main.py --video /your/video/path --model /your/model/path/best.pt得られた動画は以下のように犬の顔のみにモザイク処理が施されているでしょう。
最後に
ここまでお読みいただきありがとうございます。
ぜひ思い思いの学習データを作成して、ユニークなAIを作ってみてください。
