機械学習のデータセットのアノテーションのフォーマットはまちまちです。
望んだトレーニング用のフォーマットに変える方法をケーススタディでご紹介します。
ケーススタディ
個別ファイルのアノテーションを一つのファイルまとめ、ラベルマップを付与する
たとえば、JSONファイル一つにつき画像一枚の
・ラベル名
・ボックス
を書き込んであるアノテーションセットがあるとします。
train_000.json
{"labels": {"cat": [[1200, 700, 1800, 1300], [900, 1200, 1300, 1500], [12, 1100, 500, 1400]]}}
これを以下のように一つのファイルにまとめ、COCOデータセットのようにラベルID番号やファイル名もつけるコードがこちら。
train_annotations.json
{
"categories": [
{
"id": 1,
"name": "cat"
},
{
"id": 2,
"name": "dog"
}
],
"annotations": [
{
"filename": "train_000.jpg",
"labels": [
1,
1,
1,
],
"label_texts": [
"cat",
"cat",
"cat"
],
"boxes": [
[
[
1200,
700,
1800,
1300
],
[
900,
1200,
1300,
1500
],
[
12,
1100,
500,
1400
]
]
]
},
{
"filename": "train_001.jpg",
"labels": [
1,
1,
1,
2,
2,
2,
2,
],
"label_texts": [
"cat",
"cat",
"cat",
"dog",
"dog",
"dog",
"dog"
],
"boxes": [
[
[
1800,
500,
3000,
1200
],
[
1900,
1300,
2300,
1500
],
[
1276,
761,
1816,
1285
],
[
654,
982,
1273,
1572
]
],
[
[
1,
395,
1243,
853
],
[
1,
955,
541,
1463
]
]
},
}
import json #jsonの読み書き用ライブラリ
import os
from natsort import natsorted #ファイルのソート用
from collections import OrderedDict #順番を保った辞書を作る
import glob
json_dir = './train_annotations'
json_paths = natsorted(glob.glob(json_dir + '/*.json')) #jsonファイルたちの名前を保持
annotated_dicts = OrderedDict() #書き込み用の辞書
annotations_list = [] #書き込み用の辞書のアノテーション要素となるリスト
for json_path in json_paths: #jsonファイルたちから読み込んでいく
if os.path.exists(json_path):
with open(json_path, 'r') as f:
df = json.load(f)
filename = os.path.splitext(os.path.basename(json_path))[0] + '.jpg' #jsonファイルのファイル名（画像名に対応）を取得して画像名にする
annotated_dict = OrderedDict() #個々の画像に対するアノテーションの辞書
labels = [] #個々の画像に対するアノテーション要素の配列を作成
label_texts = []
boxes = []
for key,value in df['labels'].items():
for box in value:
label_texts.append(key) #元ファイルのラベル名をラベル名の配列へ入れる
boxes.append(box) #元ファイルのボックスをボックスの配列へ入れる
if key == 'cat':
labels.append(1) #catならID１をラベルIDの配列へ
elif key == 'dog':
labels.append(2) #dogならID２をラベルIDの配列へ
#以下で、ここまで作った要素を個々の画像に対するアノテーションの辞書に入れていく
annotated_dict['filename'] = filename
annotated_dict['labels'] = labels
annotated_dict['label_texts'] = label_texts
annotated_dict['boxes'] = boxes
annotations_list.append(annotated_dict)
annotated_dicts['categories'] = [{'id':1,'name':'cat'},
{'id':2,'name':'dog'}
] #書き込み用の辞書のカテゴリー要素となるリスト
annotated_dicts['annotations'] = annotations_list #最後に、全ての画像のアノテーション辞書を書き込み用の辞書のアノテーション要素とする
with open('./edit_train_annotations.json', 'w') as f:
json.dump(annotated_dicts, f, indent=3, ensure_ascii=False) #書き込み用の辞書をjsonファイルをとして書き込む。
🐣
フリーランスエンジニアです。
お仕事のご相談こちらまで
rockyshikoku@gmail.com
Core MLを使ったアプリを作っています。
機械学習関連の情報を発信しています。