coco-annotatorから必要な部分だけ抜き取る

何ができるコード？

coco-annotatorでkeypointだけをannotationしたデータセットのjsonからfile_nameとkeypointsだけを抜き取ってjsonファイルを作成します。
元ファイルがあるディレクトリにsimple_dataset.jsonとして出力されます。

何に使ったの？

自作keypoint detection深層学習モデルを作成中です。被写体が絶対に写っていること、全てのkeypointが見えていることを条件として簡易的なモデルを作成しています。必要な情報は画像に対応する3点のkeypointsなのでそれだけを抽出するために作成した。

コード

import json
from pathlib import Path
import sys


def open_json(path):
    with open(path) as data:
        return json.load(data)


def make_coco2simple_json(from_path, to_path):
    original_data = open_json(from_path)
    # 新しいJSONデータを作成
    new_json_data = []

    # 'images'キー内の各要素を処理
    for image in original_data["images"]:
        # ファイル名を取得
        file_name = image["file_name"]

        # アノテーション情報を取得
        keypoint = []
        for annotation in original_data["annotations"]:
            if annotation["image_id"] == image["id"]:
                keypoint = annotation["keypoints"]
                for index in range(int(len(keypoint) / 3), 0, -1):
                    del keypoint[index * 3 - 1]

        # ファイル名とアノテーション情報を新しいJSONデータに追加
        new_entry = {"file_name": file_name, "keypoints": keypoint}

        new_json_data.append(new_entry)

    # 新しいJSONデータを表示または保存
    # new_json_str = json.dumps(new_json_data, indent=4)
    # print(new_json_str)
    with open(to_path, "w") as outfile:
        json.dump(new_json_data, outfile)


if __name__ == "__main__":
    arg_len = len(sys.argv)
    input_path_string = (
        "~/Desktop/Dataset.json"
    )

    if arg_len == 1:
        pass
    elif arg_len == 2:
        input_path_string = sys.argv[1]
    else:
        print("\n予期しない引数")
        print("python json-coco2simple.py input-json-file-name")
        print("の形式で実行してください")
        exit()

    input_path = Path(input_path_string)
    if not input_path.exists():
        print(f"{input_path} は存在しません")
        exit()

    output_path = input_path.parent / "simple_dataset.json"
    # print(output_path)
    # data = open_json(input_path)
    make_coco2simple_json(from_path=input_path, to_path=output_path)

補足

cocoの仕様でkeypointには座標x,yと状態v（0:不明, 1:不可視(遮蔽), 2:可視）が含まれる。状態vは使用しないので3の倍数の部分だけ削除し、x,y座標の羅列になるようにした。