物体の三次元姿勢推定 CenterSnap -bop_toolkitでCOCO形式データ作成- 【Python】

AI
2023.12.26
スポンサーリンク
スポンサーリンク

はじめに

前回はアノテーションデータの作成を実施しました。

物体の三次元姿勢推定 CenterSnap -bop_toolkitでアノテーションデータ作成- 【Python】
物体の三次元姿勢推定 CenterSnap -bop_toolkitでアノテーションデータ作成- 【Python】 今回は、 bop_toolkit を使用してアノテーションデータを作成する方法について説明します。
nutritionfoodtech.com
2023.12.20

今回は、作成したアノテーションデータを COCO 形式に変換します。

前提条件

前提条件は以下の通りです。

アノテーションデータを COCO 形式へ変換

変換プログラムは、bop_toolkit に既にありますので、少し修正を加えて使用します。

calc_gt_coco.py

# Author: Martin Sundermeyer (martin.sundermeyer@dlr.de)
# Robotics Institute at DLR, Department of Perception and Cognition

"""Calculates Instance Mask Annotations in Coco Format."""

import numpy as np
import os
import datetime
import json

from bop_toolkit_lib import pycoco_utils
from bop_toolkit_lib import config
from bop_toolkit_lib import dataset_params
from bop_toolkit_lib import inout
from bop_toolkit_lib import misc


# PARAMETERS.
################################################################################
p = {
  # See dataset_params.py for options.
  'dataset': 'lm',

  # Dataset split. Options: 'train', 'test'.
  'dataset_split': 'train',

  # Dataset split type. Options: 'synt', 'real', None = default. See dataset_params.py for options.
  'dataset_split_type': 'pbr',

  # bbox type. Options: 'modal', 'amodal'.
  'bbox_type': 'amodal',

  # Folder containing the BOP datasets.
  'datasets_path': "/path/to/makeNOCS/output_data/bop_data",

}
################################################################################

datasets_path = p['datasets_path']
dataset_name = p['dataset']
split = p['dataset_split']
split_type = p['dataset_split_type']
bbox_type = p['bbox_type']

dp_split = dataset_params.get_split_params(datasets_path, dataset_name, split, split_type=split_type)
dp_model = dataset_params.get_model_params(datasets_path, dataset_name)

complete_split = split
if dp_split['split_type'] is not None:
    complete_split += '_' + dp_split['split_type']

CATEGORIES = [{'id': obj_id, 'name':str(obj_id), 'supercategory': dataset_name} for obj_id in dp_model['obj_ids']]
INFO = {
    "description": dataset_name + '_' + split,
    "url": "https://github.com/thodan/bop_toolkit",
    "version": "0.1.0",
    "year": datetime.date.today().year,
    "contributor": "",
    "date_created": datetime.datetime.utcnow().isoformat(' ')
}

for scene_id in dp_split['scene_ids']:
    segmentation_id = 1

    coco_scene_output = {
        "info": INFO,
        "licenses": [],
        "categories": CATEGORIES,
        "images": [],
        "annotations": []
    }

    # Load info about the GT poses (e.g. visibility) for the current scene.
    scene_gt = inout.load_scene_gt(dp_split['scene_gt_tpath'].format(scene_id=scene_id))
    scene_gt_info = inout.load_json(dp_split['scene_gt_info_tpath'].format(scene_id=scene_id), keys_to_int=True)
    # Output coco path
    coco_gt_path = dp_split['scene_gt_coco_tpath'].format(scene_id=scene_id)
    if bbox_type == 'modal':
        coco_gt_path = coco_gt_path.replace('scene_gt_coco', 'scene_gt_coco_modal')
    misc.log('Calculating Coco Annotations - dataset: {} ({}, {}), scene: {}'.format(
          p['dataset'], p['dataset_split'], p['dataset_split_type'], scene_id))
    
    # Go through each view in scene_gt
    for scene_view, inst_list in scene_gt.items():
        im_id = int(scene_view)
        
        img_path = dp_split['rgb_tpath'].format(scene_id=scene_id, im_id=im_id)
        relative_img_path = os.path.relpath(img_path, os.path.dirname(coco_gt_path))
        image_info = pycoco_utils.create_image_info(im_id, relative_img_path, dp_split['im_size'])
        coco_scene_output["images"].append(image_info)
        gt_info = scene_gt_info[scene_view]
        
        # Go through each instance in view
        for idx,inst in enumerate(inst_list): 
            category_info = inst['obj_id']
            visibility = gt_info[idx]['visib_fract']
            # Add ignore flag for objects smaller than 10% visible
            ignore_gt = visibility < 0.1
            mask_visib_p = dp_split['mask_visib_tpath'].format(scene_id=scene_id, im_id=im_id, gt_id=idx)
            mask_full_p = dp_split['mask_tpath'].format(scene_id=scene_id, im_id=im_id, gt_id=idx)
            
            binary_inst_mask_visib = inout.load_depth(mask_visib_p).astype(np.bool)
            if binary_inst_mask_visib.sum() < 1:
                continue
            if bbox_type == 'amodal':
                binary_inst_mask_full = inout.load_depth(mask_full_p).astype(np.bool)
                if binary_inst_mask_full.sum() < 1:
                    continue
                bounding_box = pycoco_utils.bbox_from_binary_mask(binary_inst_mask_full)
            elif bbox_type == 'modal':
                bounding_box = pycoco_utils.bbox_from_binary_mask(binary_inst_mask_visib)
            else:
                raise Exception('{} is not a valid bounding box type'.format(p['bbox_type']))

            annotation_info = pycoco_utils.create_annotation_info(
                segmentation_id, im_id, category_info, binary_inst_mask_visib, bounding_box, tolerance=2, ignore=ignore_gt)

            if annotation_info is not None:
                coco_scene_output["annotations"].append(annotation_info)

            segmentation_id = segmentation_id + 1

    with open(coco_gt_path, 'w') as output_json_file:
        json.dump(coco_scene_output, output_json_file)

上記を実行すると、scene_gt_coco.json が作成されます。

cd makeNOCS/bop_toolkit
python3 scripts/calc_gt_coco.py

scene_gt_coco.json

{"info": {"description": "lm_train", 
          "url": "https://github.com/thodan/bop_toolkit", 
          "version": "0.1.0", 
          "year": 2023, 
          "contributor": "", 
          "date_created": "2023-05-15 06:32:14.908762"
         }, 
 "licenses": [], 
 "categories": [{"id": 1, "name": "1", "supercategory": "lm"}], 
 "images": [{"id": 0, "file_name": "rgb/000000.jpg", "width": 640, "height": 480, "date_captured": "2023-05-15 06:32:14.917984", "license": 1, "coco_url": "", "flickr_url": ""}, 
            {"id": 1, "file_name": "rgb/000001.jpg", "width": 640, "height": 480, "date_captured": "2023-05-15 06:32:15.171713", "license": 1, "coco_url": "", "flickr_url": ""}, 
　　　　　　 {"id": 2, "file_name": "rgb/000002.jpg", "width": 640, "height": 480, "date_captured": "2023-05-15 06:32:15.396059", "license": 1, "coco_url": "", "flickr_url": ""},

COCO 形式の見たことあるような構成になっています。

プログラム説明

import numpy as np
import os
import datetime
import json

from bop_toolkit_lib import pycoco_utils
from bop_toolkit_lib import config
from bop_toolkit_lib import dataset_params
from bop_toolkit_lib import inout
from bop_toolkit_lib import misc


# PARAMETERS.
################################################################################
p = {
  # See dataset_params.py for options.
  'dataset': 'lm',

  # Dataset split. Options: 'train', 'test'.
  'dataset_split': 'train',

  # Dataset split type. Options: 'synt', 'real', None = default. See dataset_params.py for options.
  'dataset_split_type': 'pbr',

  # bbox type. Options: 'modal', 'amodal'.
  'bbox_type': 'amodal',

  # Folder containing the BOP datasets.
  'datasets_path': "/path/to/makeNOCS/output_data/bop_data",

}
################################################################################

今までと異なるのは、bbox_type です。

  • modal … 遮蔽された部分を無視して bbox を作成
  • amodal … 遮蔽された部分も加味して bbox を作成

今回は物体の中心も知りたいので、amodal とします。

datasets_path = p['datasets_path']
dataset_name = p['dataset']
split = p['dataset_split']
split_type = p['dataset_split_type']
bbox_type = p['bbox_type']

dp_split = dataset_params.get_split_params(datasets_path, dataset_name, split, split_type=split_type)
dp_model = dataset_params.get_model_params(datasets_path, dataset_name)

complete_split = split
if dp_split['split_type'] is not None:
    complete_split += '_' + dp_split['split_type']

この辺も今までと同じです。

dataset_params.py からデータを取得しています。

CATEGORIES = [{'id': obj_id, 'name':str(obj_id), 'supercategory': dataset_name} for obj_id in dp_model['obj_ids']]
INFO = {
    "description": dataset_name + '_' + split,
    "url": "https://github.com/thodan/bop_toolkit",
    "version": "0.1.0",
    "year": datetime.date.today().year,
    "contributor": "",
    "date_created": datetime.datetime.utcnow().isoformat(' ')
}

COCO データとして出力するための初期情報を取得します。

for scene_id in dp_split['scene_ids']:
    segmentation_id = 1

    coco_scene_output = {
        "info": INFO,
        "licenses": [],
        "categories": CATEGORIES,
        "images": [],
        "annotations": []
    }

各シーンの情報を揃えていきます。
segmentation_id は後述します。

# Load info about the GT poses (e.g. visibility) for the current scene.
scene_gt = inout.load_scene_gt(dp_split['scene_gt_tpath'].format(scene_id=scene_id))
scene_gt_info = inout.load_json(dp_split['scene_gt_info_tpath'].format(scene_id=scene_id), keys_to_int=True)
# Output coco path
coco_gt_path = dp_split['scene_gt_coco_tpath'].format(scene_id=scene_id)
if bbox_type == 'modal':
    coco_gt_path = coco_gt_path.replace('scene_gt_coco', 'scene_gt_coco_modal')
misc.log('Calculating Coco Annotations - dataset: {} ({}, {}), scene: {}'.format(
            p['dataset'], p['dataset_split'], p['dataset_split_type'], scene_id))

scene_gt は scene_gt.json、scene_gt_info は、scene_gt_info.json を読込みます。

coco_gt_path は出力するファイル名を指定しますが、デフォルトで問題ありません。
bbox_type を modal にした場合は、ファイル名が変わります。

# Go through each view in scene_gt
for scene_view, inst_list in scene_gt.items():
    im_id = int(scene_view)

scene_gt.json の情報を順次処理します。

img_path = dp_split['rgb_tpath'].format(scene_id=scene_id, im_id=im_id)
relative_img_path = os.path.relpath(img_path, os.path.dirname(coco_gt_path))
image_info = pycoco_utils.create_image_info(im_id, relative_img_path, dp_split['im_size'])
coco_scene_output["images"].append(image_info)
gt_info = scene_gt_info[scene_view]

画像と scene_gt_info.json の情報を読込みます。

# Go through each instance in view
for idx,inst in enumerate(inst_list): 
    category_info = inst['obj_id']
    visibility = gt_info[idx]['visib_fract']
    # Add ignore flag for objects smaller than 10% visible
    ignore_gt = visibility < 0.1
    mask_visib_p = dp_split['mask_visib_tpath'].format(scene_id=scene_id, im_id=im_id, gt_id=idx)
    mask_full_p = dp_split['mask_tpath'].format(scene_id=scene_id, im_id=im_id, gt_id=idx)

一枚の画像内のインスタンス毎にループを回していきます。

カテゴリ、可視度(視認度？)、可視部分のマスク画像、マスク画像を読込みます。

binary_inst_mask_visib = inout.load_depth(mask_visib_p).astype(np.bool)
if binary_inst_mask_visib.sum() < 1:
    continue
if bbox_type == 'amodal':
    binary_inst_mask_full = inout.load_depth(mask_full_p).astype(np.bool)
    if binary_inst_mask_full.sum() < 1:
        continue
    bounding_box = pycoco_utils.bbox_from_binary_mask(binary_inst_mask_full)
elif bbox_type == 'modal':
    bounding_box = pycoco_utils.bbox_from_binary_mask(binary_inst_mask_visib)
else:
    raise Exception('{} is not a valid bounding box type'.format(p['bbox_type']))

可視部分のマスク画像を読込んで処理していきます。

オブジェクトがすべて隠れている場合は、パスします。
今回は amodal な bbox なので、オブジェクト全体のマスク画像から、bbox 情報を読込みます。

annotation_info = pycoco_utils.create_annotation_info(
       segmentation_id, im_id, category_info, binary_inst_mask_visib, bounding_box, tolerance=2, ignore=ignore_gt)

if annotation_info is not None:
    coco_scene_output["annotations"].append(annotation_info)

segmentation_id = segmentation_id + 1

最後に COCO アノテーション情報を作成していきます。

  • segmentation_id … インスタンスごとの重複無しの整数ID
  • image_id … 画像ごとの重複無しの整数ID
  • object_id … カテゴリに登録された番号、今回は 1 のみ
  • binary_mask … バイナリ形式のマスクデータ
  • bouding_box … [x, y, w, h] の bbox
  • mask_encoding_format … デフォルトでは rle
  • tolerance … ポリゴンをマスクにフィットさせるための許容誤差
  • ignore … このアノテーションを無視するかどうか
with open(coco_gt_path, 'w') as output_json_file:
    json.dump(coco_scene_output, output_json_file)

最後に、json データを出力して終了です。

おわりに

今回はCOCO形式のアノテーションデータを作成しました。

次回は、オブジェクトの情報を作成していきます。

また、現在シミュレーションデータから実データに置き換える部分を対応中ですが、もしかしたら上手くいかないかもしれません。

物体の三次元姿勢推定 CenterSnap -オブジェクト情報の作成- 【Python】
物体の三次元姿勢推定 CenterSnap -オブジェクト情報の作成- 【Python】 今回は、オブジェクト情報を作成していきます。
nutritionfoodtech.com
2023.12.30
スポンサーリンク

コメント

タイトルとURLをコピーしました