YOLO 特定のクラスだけ検出したい

前提・実現したいこと

yolov2を用いて精度の良い物体検出をしたいと考えています。
検出する物をcarのみにしたい為、class_namesを変更など試しているのですが指定した物以外も検出されてしまいます。
なにかやり方があれば教えていただきたいです。

該当のソースコード

from darkflow.net.build import TFNet
import cv2
import numpy as np

options = {"model": "yolov2-tiny-voc.cfg", "load": "yolov2-tiny-voc.weights", "threshold": 0.1,"GPU": 1.0}
tfnet = TFNet(options)

# カメラの起動
cap = cv2.VideoCapture(r"C:\Script\_home_pi_movie__home_pi_movie_cap20201004_1411.mp4")

class_names = ["aeroplane", "bicycle", "bird", "boat", "bottle",
    "bus", "car", "cat", "chair", "cow", "diningtable", "dog",
    "horse", "motorbike", "person", "pottedplant", "sheep", "sofa",
    "train", "tvmonitor"]

width = int(cap.get(3))
height = int(cap.get(4))
fourcc = cv2.VideoWriter_fourcc('H','2','6','4')
out = cv2.VideoWriter('tesaaat.mp4',fourcc, 15, (width,height))

num_classes = len(class_names)
class_colors = []
for i in range(0, num_classes):
    hue = 255*i/num_classes
    col = np.zeros((1,1,3)).astype("uint8")
    col[0][0][0] = hue
    col[0][0][1] = 128
    col[0][0][2] = 255
    cvcol = cv2.cvtColor(col, cv2.COLOR_HSV2BGR)
    col = (int(cvcol[0][0][0]), int(cvcol[0][0][1]), int(cvcol[0][0][2]))
    class_colors.append(col) 
    
def main():
    
    while(True):
        
        # 動画ストリームからフレームを取得
        ret, frame = cap.read()
        result = tfnet.return_predict(frame)
        #cv2.line(frame, (0, pos_linha), (1280, pos_linha), (255,127,0), 1) #線

        for item in result:
            tlx = item['topleft']['x']
            tly = item['topleft']['y']
            brx = item['bottomright']['x']
            bry = item['bottomright']['y']
            label = item['label']
            conf = item['confidence']

            if conf > 0.5:
                
                for i in class_names:
                    if label == i and class_names[7] == label:
                        class_num = class_names.index(i)
                        break
                
                #枠の作成
                cv2.rectangle(frame, (tlx, tly), (brx, bry), class_colors[class_num], 2)
                #ラベルの作成
                text = label + " " + ('%.2f' % conf)  
                cv2.rectangle(frame, (tlx, tly - 15), (tlx + 100, tly + 5), class_colors[class_num], -1)
                cv2.putText(frame, text, (tlx, tly), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,0,0), 1)

        # 表示
        cv2.imshow("Show FLAME Image", frame)
        out.write(frame)

        # escを押したら終了。
        k = cv2.waitKey(10);
        if k == ord('q'):  break;

    cap.release()
    out.release()
    cv2.destroyAllWindows()
    
if __name__ == '__main__':
    main()

行動規範の内容に同意します

回答2件

↓の部分が何をしたいか意図が分かりませんが、頑張ってcarだけにしようと試みてる部分でしょうか？

                for i in class_names:
                    if label == i and class_names[7] == label:
                        class_num = class_names.index(i)
                        break

temps1101さんが言われている通り、本来はクラス数を限定するべきですが、とりあえず表示だけそれっぽく動かすなら↓みたいに誤魔化せますけどね。

            if conf > 0.5 and label=="car":
                #枠の作成
                cv2.rectangle(frame, (tlx, tly), (brx, bry), class_colors[class_num], 2)
                #ラベルの作成
                text = label + " " + ('%.2f' % conf)  
                cv2.rectangle(frame, (tlx, tly - 15), (tlx + 100, tly + 5), class_colors[class_num], -1)
                cv2.putText(frame, text, (tlx, tly), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,0,0), 1)

投稿2020/11/19 14:45

編集2020/11/19 14:46