YOLOv7でバウンディングボックスの色を指定する方法をメモしておきます。
YOLOv7では、バウンディングボックスの色がランダムになるため、この方法をメモすることにしました。
Google Colaboratoryにて実行しました。
ランタイム->ランタイムのタイプを変更->T4 GPUを選択
from google.colab import drive
drive.mount('/content/drive')
%cd "/content/drive/MyDrive"
!git clone https://github.com/WongKinYiu/yolov7
%cd "/content/drive/MyDrive/yolov7"
!pip install -r requirements.txt
物体検出(結果はyolov7/runs/detectに保存される)
%cd "/content/drive/MyDrive/yolov7"
!python detect.py --weights yolov7.pt --conf 0.25 --img-size 640 --source "/content/drive/MyDrive/yolov7/inference/images/horses.jpg"
この画像が出力されます。
もう一回実行します。
バウンディングボックスの色がランダムになることがわかります。
そこで、バウンディングボックスの色を指定するため、detect.pyの62行目に以下のコードを挿入します。
色の順番はBGRです。
今回は黄緑色に指定しました。
detect.py
# 62行目
colors = [[20, 255, 181] for _ in names] # 色の順番はBGR
以下のコードをGoogle Colaboratoryにて再度実行します。
%cd "/content/drive/MyDrive/yolov7"
!python detect.py --weights yolov7.pt --conf 0.25 --img-size 640 --source "/content/drive/MyDrive/yolov7/inference/images/horses.jpg"
detect.pyを書き換えることで、バウンディングボックスの色を指定することができました。
この方法では、全てのクラスのバウンディングボックスの色がこの色に指定されてしまうため、クラスごとに色を変えられるように改善する必要がありそうです。