label_map.pbtxt をつくる

Tensorflow Object Detection API に必要なlabel_map。
手作業で書き換えるのは面倒ということで、ラベルの配列から作れるコードがこちら。

label_map_path = "./test.pbtxt" #保存パス
categories = ["dog","cat","bird"] #ラベルの配列

end = '\n'
s = ' '
class_map = {}
for ID, name in enumerate(categories):
    out = ''
    out += 'item' + s + '{' + end
    out += s*2 + 'id:' + ' ' + (str(ID+1)) + end
    out += s*2 + 'name:' + ' ' + '\'' + name + '\'' + end
    out += '}' + end*2
    

    with open(label_map_path, 'a') as f:
        f.write(out)
        
    class_map[name] = ID+1

これで出力保存されます。

label_map.pbtxt

item {
  id: 1
  name: 'dog'
}

item {
  id: 2
  name: 'cat'
}

item {
  id: 3
  name: 'bird'
}

このlabel_map.pbtxtで物体検出のラベルIDとラベル名を対応づけられます。

ObjectDetectionAPIの使い方はこちら

🐣

---

フリーランスエンジニアです。
お仕事のご相談こちらまで
[e-mail]

Core MLを使ったアプリを作っています。
機械学習関連の情報を発信しています。

Twitter
Medium