会社のファミリーデー(職場見学イベント)に子どもを招待したので、見てもらうモノのひとつとして表題のAIを実装してみました。
環境
- 学習環境
Google Colaboratory(無料枠) + Google Drive(要400MB程度の空き容量) - 推論環境
Windows10(i7-10750H/16GB/RTX2060 Max-Q/内臓カメラ1280x720)
VSCode
Python 3.9.13 - モデル
YOLOv8
最新のYOLO-NASがありますが、GitHubにTrainコードがなく、実装するのには時間がなかったため、サンプルがすぐ見つかったYOLOv8を採用しました。
データセット
ありがたいことに、YOLO形式のアノテーションファイルが同梱されたZIPファイルをダウンロードすることができます。
上記のデータセットは、以下の量が入っています。
- train: 1366画像
- valid: 368画像
- test: 200画像
Windows10エクスプローラーのZIP解凍だと2ファイルほどエラーになるかもしれません。（ファイル名が長すぎる？）
7-Zipファイルマネージャーだと正常に解凍できました。
解凍後は206MBになります。
画像を見ると、マスクあり/なしの人以外にも背景だけの画像もあり、色々と工夫されているように感じます。
解凍後のフォルダ一式をGoogleDriveにアップロードします。
(GoogleDriveのROOT直下に「mask」などのフォルダを作成してアップロードすると、後のcdコマンドを打つのが楽になる)
学習
先ほどデータセットをアップロードしたGoogleDrive上で[新規-その他-Google Colaboratory]を選択し、Goole Colab画面に遷移します。
まずは、Python3とGPUを使うようランタイムを設定します。
メニューの[ランタイム-ランタイムのタイプを変更]をクリックし、[Python3]と[T4 GPU]にして保存します。
次にGoogle Driveをマウントします。
from google.colab import drive
drive.mount('/content/drive')
%cd /content/drive/MyDrive
%ll
#先ほどアップロードしたフォルダ名があればOK
test/valid/testフォルダがあるところに移動します。
%cd mask
%ll
#data.yaml/train/valid/testがあればOK
YOLOv8をダウンロードします。
pipかGitHubクローンの2通りありますが、pip推奨とのことなので、以下のpipを実行します。
!pip3 install ultralytics
#「Successfully installed ultralytics-8.0.145」のように表示されればOK
学習を実行します。
!yolo task=detect mode=train model=yolov8n.pt data=data.yaml epochs=300 imgsz=640 batch=64
モデルは以下があります。
| モデル | 性能 |
|---|---|
| YOLOv8n | 軽量、精度低い |
| YOLOv8s | |
| YOLOv8m | 中間 |
| YOLOv8l | |
| YOLOv8x | 重い、精度高い |
モデルの選択
推論の速度と精度のバランスを見て、適したものを選択する必要があります。
今回は普通のノートPCでリアルタイム推論するため、一番軽量のYOLOv8nを選択しました。
今回の推論環境だと、YOLOv8nだと推論1回に200ms程度かかりパラパラと描画されましたが、YOLOv8xだと4000msもかかるためハングアップしているかのようになり使い物になりませんでした。
学習中の注意事項
学習には2時間程度かかります。
何も操作していない時間が90分続くと切断されるため、たまにブラウザ上でマウスを動かす必要があります。
また、GPUの使用制限(詳細非公開)があるため、「これ以上実行する場合はProを購入してね」と出ることがあります。
（Google側の使用リソースに依るようで、そのまま学習完了することもありました）
もし使用制限になってしまった場合、翌日に以下のコマンドで再開することができます。
!yolo task=detect mode=train resume model='./runs/detect/train/weights/last.pt' data=data.yaml epochs=300 imgsz=640
評価
学習が終わると「./runs/detect/train」以下に結果が出力されています。
この中に「results.png」を開くと以下のようになっていました。
- lossは、正解と予測結果との差（小さいほどbetter)
- precisionは、"確実にマスクを着けている人"と予測して的中した結果値
- recallは、"外れているかもしれないけど、とにかくマスクを着けているだろう人"と予測して的中した結果値
今回は、lossが小さくなって推移しており、precision/recallも高く推移しているので精度は問題ないだろうと判断しました。
推論(実験)
「./runs/detect/train/weights/best.pt」を推論環境(Windows)にダウンロードします。
(YOLOv8nは6MBでした)
ptファイルは"重みファイル"と呼ばれ、YOLOv8nモデル内の重み(係数)が記録されたファイルです。
推論環境(Windows)でも「pip3 install ultralytics」を実行してYOLOv8をインストールしておきます。
VSCodeで以下のコードを実行してみます。
from ultralytics import YOLO
model = YOLO("best.pt")
results = model(0 , show=True)
for i in enumerate(results):
print(i)
コードはコチラを参考にしました。
https://zenn.dev/opamp/articles/51ee26445a1732#show%E3%81%AE%E6%89%B1%E3%81%84
model(0, ...)と指定すると、データソースがWebカメラになるようです。
実行して30秒ほど待つと、カメラ映像のウィンドウが表示され、人の顔を認識してマスクあり/なしを表示してくれます。
VSCodeのコンソールには、推論結果+処理時間(ms単位)が表示されます。
これで性能が問題ないか確認していきます。
- マスクをつけてmaskと判定されるか
- マスクを外してno-maskと判定されるか
- カメラ映像がカクカクしすぎていないか
性能は問題ないことが確認できたら、以下の問題を解消していきます。
- maskとno-maskの枠の色が似ている
- カメラ映像が左右反転していて不自然
推論(本番)
VSCodeで以下のコードを実行します。
コードはコチラを参考にしました。
もっと最小限でWebカメラリアルタイム推論がしたい
from ultralytics import YOLO
import torch
import cv2
from ultralytics.yolo.data.augment import LetterBox
from ultralytics.yolo.utils.plotting import Annotator, colors
from ultralytics.yolo.utils import ops
from copy import deepcopy
import numpy as np
import matplotlib.pyplot as plt
model = YOLO("best.pt")
cap = cv2.VideoCapture(0)
def preprocess(img, size=1280):
img = LetterBox(size, True)(image=img)
img = img.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB
img = np.ascontiguousarray(img) # contiguous
img = torch.from_numpy(img)
img = img.float() # uint8 to fp16/32
img /= 255 # 0 - 255 to 0.0 - 1.0
return img.unsqueeze(0)
def postprocess(preds, img, orig_img):
preds = ops.non_max_suppression(preds,
0.25,
0.8,
agnostic=False,
max_det=100)
for i, pred in enumerate(preds):
shape = orig_img.shape
pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], shape).round()
return preds
def draw_bbox(pred, names, annotator):
for *xyxy, conf, cls in reversed(pred):
c = int(cls) # integer class
label = f'{names[c]} {conf:.2f}'
annotator.box_label(xyxy, label, color=colors(c*5, True))
while True:
ret, img = cap.read()
img = cv2.flip(img, 1) #追加
origin = deepcopy(img)
annotator = Annotator(origin,line_width=1,example=str(model.model.names))
img = preprocess(img)
preds = model.model(img, augment=False)
preds = postprocess(preds,img,origin)
draw_bbox(preds[0], model.model.names, annotator)
cv2.imshow("test",origin)
cv2.waitKey(1)
「maskとno-maskの枠の色が似ている」問題の対策
draw_bbox関数のcolor=colors(c, True)を変更しました。
colorsは以下で定義されていました。
class Colors:
"""Ultralytics color palette https://ultralytics.com/."""
def __init__(self):
"""Initialize colors as hex = matplotlib.colors.TABLEAU_COLORS.values()."""
hexs = ('FF3838', 'FF9D97', 'FF701F', 'FFB21D', 'CFD231', '48F90A', '92CC17', '3DDB86', '1A9334', '00D4BB',
'2C99A8', '00C2FF', '344593', '6473FF', '0018EC', '8438FF', '520085', 'CB38FF', 'FF95C8', 'FF37C7')
self.palette = [self.hex2rgb(f'#{c}') for c in hexs]
self.n = len(self.palette)
self.pose_palette = np.array([[255, 128, 0], [255, 153, 51], [255, 178, 102], [230, 230, 0], [255, 153, 255],
[153, 204, 255], [255, 102, 255], [255, 51, 255], [102, 178, 255], [51, 153, 255],
[255, 153, 153], [255, 102, 102], [255, 51, 51], [153, 255, 153], [102, 255, 102],
[51, 255, 51], [0, 255, 0], [0, 0, 255], [255, 0, 0], [255, 255, 255]],
dtype=np.uint8)
カメラ映像が左右反転していて不自然
カメラからキャプチャした画像をcv2.flip関数で左右反転するだけです。
while True:
ret, img = cap.read()
img = cv2.flip(img, 1) #追加
まとめ
出来合いのデータセットと推論コードで簡単にマスク付けてる・付けていないを判定するAIを作ることができました。
ファミリーデー当日、４～５人が一度に映り込んでもそれぞれ判定できていました。
鼻出しマスク状態の判定が怪しそうでした。
おそらく学習データセットの中に、そういうデータ＆クラスが無いからだと思われます。