機械学習のデータセットのアノテーションのフォーマットはまちまちです。
望んだトレーニング用のフォーマットに変える方法をケーススタディでご紹介します。

#ケーススタディ

###個別ファイルのアノテーションを一つのファイルまとめ、ラベルマップを付与する

たとえば、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を使ったアプリを作っています。
機械学習関連の情報を発信しています。

Twitter
Medium