【YOLOv8】BBOXの色の変更方法

YOLOv8のバウンディングボックスの色を変更

YOLOv8を用いてヘルメットやハーネスなどを着用しているかを自動でチェックするシステムの開発を行っている。

デモを行った際、BBOXの文字の色が分かりにくい、見えずらいなどの意見が出たため変更の仕方のメモを残す。

以下がデモで使用したコード

import cv2
from ultralytics import YOLO
model_chin_detect = YOLO(学習済みの重み) 

cap = cv2.VideoCapture(1)#ウェブカメラ

while cap.isOpened():
    success, frame = cap.read()
    k = cv2.waitKey(1) # 一応キー入力で終了できるようにしておく

    if k != -1:
        break
    if success:
        chin_results = model_chin_detect(frame,conf=0.5)
        annotated_frame = chin_results[0].plot()
        # clothes_results = model_detect(annotated_frame,conf=0.9)
        # annotated_frame = clothes_results[0].plot()
        cv2.imshow("YOLOv8 Inference", annotated_frame)
cap.release()
cv2.destroyAllWindows()

動作環境

• Windows11
• ultralytics == 8.0.58

修正箇所を探す

まずgithubからYOLOv8のソースコードを落としてくる。

git clone https://github.com/ultralytics/ultralytics

YOLOv8のdetectの処理を追っているとresults.pyのResultsクラスの228行目に描画系の処理っぽいのを発見。

annotator = Annotator(
            deepcopy(self.orig_img if img is None else img),
            line_width,
            font_size,
            font,
            pil or (pred_probs is not None and show_probs),  # Classify tasks default to pil=True
            example=names,
        )

この処理はresults.pyの258行目のannotator.box_label()で使用されている。

plotting.pyからimportされているのでplotting.pyのannotatorクラスを見に行く。

163行目付近のbox_labelメソッドにopenCVのcv2.rectangleが使われているためここらへんでバウンディングボックスの描画を行っているっぽい？？

163行目のbox_labelメソッドの引数に検出したラベルの情報を渡してラベル名に応じてif文でBBoxの色を変更できそう。

このメソッドの引数labelにその情報が入っているがresults.pyで呼び出した際にラベル名や検出確率を文字列で連結して渡しているため不要な情報が含まれている。なのでlabel_onlyを引数に追加。

#デフォルトの引数
def box_label(self, box, label="",color=(128, 128, 128), txt_color=(255, 255, 255), rotated=False):

#変更後の引数
def box_label(self, box, label="",label_only="",color=(128, 128, 128), txt_color=(255, 255, 255), rotated=False):

label_onlyを追加したのでresults.pyのresultsクラスで呼び出されているbox_labelにも変更を加える。

ultralytics/engin/results.pyの251行目から259行目抜粋

# Plot Detect results
        if pred_boxes is not None and show_boxes:
            for d in reversed(pred_boxes):
                c, conf, id = int(d.cls), float(d.conf) if conf else None, None if d.id is None else int(d.id.item())
                name = ("" if id is None else f"id:{id} ") + names[c]
                label = (f"{name} {conf:.2f}" if conf else name) if labels else None
                label_only = f"{name}"
                box = d.xyxyxyxy.reshape(-1, 4, 2).squeeze() if is_obb else d.xyxy.squeeze()
                annotator.box_label(box, label,label_only, color=colors(c, True), rotated=is_obb)

lable情報だけを取得できたのでlabelの種類で条件分岐し色の変更を行う。

ultralytics/utils/plotting.pyの184行目から213行目抜粋　変更前

else:  # cv2
            if rotated:
                p1 = [int(b) for b in box[0]]
                # NOTE: cv2-version polylines needs np.asarray type.
                cv2.polylines(self.im, [np.asarray(box, dtype=int)], True, color, self.lw)
            else:
                p1, p2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
                cv2.rectangle(self.im, p1, p2, color, thickness=self.lw, lineType=cv2.LINE_AA)
            if label:
                print("into if")
                w, h = cv2.getTextSize(label, 0, fontScale=self.sf, thickness=self.tf)[0]  # text width, height
                outside = p1[1] - h >= 3
                p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3
                cv2.rectangle(self.im, p1, p2, color, -1, cv2.LINE_AA)  # filled
                cv2.putText(
                    self.im,
                    label,
                    (p1[0], p1[1] - 2 if outside else p1[1] + h + 2),
                    0,
                    self.sf,
                    txt_color,
                    thickness=self.tf,
                    lineType=cv2.LINE_AA,
                )

ultralytics/utils/plotting.pyの184行目から213行目抜粋　変更後

else:  # cv2
            if rotated:
                p1 = [int(b) for b in box[0]]
                # NOTE: cv2-version polylines needs np.asarray type.
                cv2.polylines(self.im, [np.asarray(box, dtype=int)], True, color, self.lw)
            else:
                print("into if")
                p1, p2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
                if label_only == "buckle is open":
                    cv2.rectangle(self.im, p1, p2, (0,0,255), thickness=self.lw, lineType=cv2.LINE_AA)
                elif label_only == "above the chin":
                    cv2.rectangle(self.im, p1, p2, (0,0,255), thickness=self.lw, lineType=cv2.LINE_AA)
                elif label_only == "right":
                    cv2.rectangle(self.im, p1, p2, (255,0,0), thickness=self.lw, lineType=cv2.LINE_AA)
            if label:
                w, h = cv2.getTextSize(label, 0, fontScale=self.sf, thickness=self.tf)[0]  # text width, height
                outside = p1[1] - h >= 3
                p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3
                # cv2.rectangle(self.im, p1, p2, color, -1, cv2.LINE_AA)  # filled
                if label_only == "buckle is open":
                    cv2.rectangle(self.im, p1, p2, (0,0,255), -1, cv2.LINE_AA)
                elif label_only == "above the chin":
                    cv2.rectangle(self.im, p1, p2, (0,0,255), -1, cv2.LINE_AA)
                elif label_only == "right":
                    cv2.rectangle(self.im, p1, p2, (255,0,0), -1, cv2.LINE_AA)
                cv2.putText(
                    self.im,
                    label,
                    (p1[0], p1[1] - 2 if outside else p1[1] + h + 2),
                    0,
                    self.sf,
                    txt_color,
                    thickness=self.tf,
                    lineType=cv2.LINE_AA,
                )

以上でバウンディングボックスの色を変更できました。