yolov3の検知種類の制限と表示方法の変更

yolov3で物体検知を行っているのですが80種類あるデータのうち制限をして、例えば人のみ検知するように設定することは可能でしょうか？
また、表示方法に関して枠ではなく検知部分の中央もしくは下部にポイントを打つような表示に変更することができますでしょうか？
素人質問で申し訳ないです．．．

anacondaで環境設定をしており、言語はpythonを使用しています。(Mac)
必要なモジュール等はインストール済みです。

tiitoi

2021/02/17 07:25

yolo といってもたくさん実装があるので、どのコードなのかわかりませんが、描画部分のソースコードを読んで、「MSCOCOの場合クラスID0が人なので、それ以外のクラスを無視する」「中央に点を描画する」の2点を修正すればいいと思います。

tetetei

2021/02/17 07:32

ご回答ありがとうございます。 https://qiita.com/daiarg/items/03c55c82fdc6e623bf07 このサイトのコードを使用しています。 writeの関数のところの書き換えだとはおもうのですが、if文を使って書くのでしょうか？

tiitoi

2021/02/17 07:38

コード自体は動かせたのでしょうか？ def write(x, img) で矩形の描画を行っているようなので、その部分のみでよければ回答しますが

tetetei

2021/02/17 07:41

コード自体は動かせています！色などの変更等もうまく対応できました。描画と人のみにするところを教えていただけるとありがたいです。今矩形で表示している部分の下の辺の中央部にポイントを置くような表示を目指しています。お手数おかけします。

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

回答1件

ベストアンサー

動作確認はしていませんが、これでどうでしょうか。

python
1def write(x, img):
2    c1 = tuple(x[1:3].int())  # 左上の座標 (x, y)
3    c2 = tuple(x[3:5].int())  # 右下の座標 (x, y)
4
5    if int(x[-1]) != 0:
6        return  # ID=0 が人なので、人以外はなにもしない
7    
8    # 矩形の中心点
9    cx = (c1[0] + c2[0]) // 2  # (左上の x 座標 + 右下の x 座標) / 2
10    cy = (c1[1] + c2[1]) // 2  # (左上の y 座標 + 右下の y 座標) / 2
11
12    # 点 (cx, cy) に半径2の円を赤色(0, 0, 255)で描画する。
13    cv2.circle(img, (cx, cy), radius=2, color=(0, 0, 255), thickness=-1)
14
15    return img