自分のデータでCOCO形式のデータセットを正しく作るの、本当にこれであっているのかなあ、と不安になりながらやっていたので、
これでOKだよ、というのをメモ。
ちなみにObject Detectionタスクだけに使うことを念頭に置いています。
Segmentationタスクはやったことない。
Object Detection, Segmentationのデータの作り方
こちらの記事のコードを少しだけ改変
import json
import collections as cl
import cv2
def get_info():
tmp = cl.OrderedDict()
tmp["description"] = "my own dataset"
tmp["url"] = ""
tmp["version"] = "1"
tmp["year"] = 2023
tmp["contributor"] = ""
tmp["data_created"] = "2023/02/09"
return tmp
def get_licenses():
tmp = cl.OrderedDict()
tmp["id"] = 0
tmp["url"] = ""
tmp["name"] = ""
return tmp
def get_image_data(id, image_path, image_file, w, h):
tmp = cl.OrderedDict()
tmp["license"] = 0
tmp["id"] = id
tmp["file_name"] = image_file
tmp["width"] = w
tmp["height"] = h
tmp["date_captured"] = ""
tmp["coco_url"] = ""
tmp["flickr_url"] = ""
return tmp
def get_annotation(id, image_id, category_id, bbox, segmentation):
tmp = cl.OrderedDict()
tmp["segmentation"] = segmentation
tmp["id"] = id
tmp["image_id"] = image_id
tmp["category_id"] = category_id
tmp["area"] = bbox[2] * bbox[3]
tmp["iscrowd"] = 0
tmp["bbox"] = bbox
return tmp
def get_categories():
tmps = []
sup = ["my super category"]
categories = ["category0", "category1", "category2"]
for i in range(len(defect)):
tmp = cl.OrderedDict()
tmp["id"] = i
tmp["supercategory"] = sup
tmp["name"] = categories[i]
tmps.append(tmp)
return tmps
以下は例えば、Yolov5で使う画像ごとのアノテーションテキストファイルをCOCOに変換する例。
image_dir = "my_dataset/images/train/"
label_dir = "my_dataset/labels/train/"
json_path = "my_dataset_train.json"
image_list = os.listdir(image_dir)
print(len(image_list))
info = get_info()
licenses = get_licenses()
categories = get_categories()
images = []
annotations = []
for i, image_file in enumerate(image_list):
image_file = image_list[i]
image_path = os.path.join(image_dir,image_file)
label_file = os.path.splitext(os.path.basename(image_file))[0] + ".txt"
label_path = os.path.join(label_dir,label_file)
if os.path.exists(label_path):
img = cv2.imread(image_path)
img_h, img_w, _ = img.shape
image_data = get_image_data(i, image_path, image_file, img_w, img_h)
source_file = open(label_path)
for object_index, line in enumerate(source_file):
staff = line.split()
class_idx = int(staff[0])
x_center, y_center, width, height = float(staff[1])*img_w, float(staff[2])*img_h, float(staff[3])*img_w, float(staff[4])*img_h
x = round(x_center-width/2,2)
y = round(y_center-height/2,2)
width = round(width,2)
height = round(height,2)
bbox = [x, y, width, height]
id = i * 1000 + object_index
annotation = get_annotation(id, i, class_idx, bbox, [])
annotations.append(annotation)
images.append(image_data)
json_data = {
'info': info,
'images': images,
'licenses': licenses,
'annotations': annotations,
'categories': categories,
}
with open(json_path, 'w', encoding='utf-8') as f:
json.dump(json_data, f, ensure_ascii=False)
作ったデータセットを視覚化して確認する。
COCO-Viewerで視覚化できる。
使い方
ローカルにクローン（Colabでは使えないので自分のパソコンに）。
python3 cocoviewer.py -i my_dataset/images/train -a my_dataset_train.json
-i に画像ディレクトリ
-a にjsonファイルを指定する。
[COCO-Viewerのリポジトリより画像]
うまく表示されなければ、データセットjsonの作り方がどこか間違っているので、公式のCOCOデータセットをダウンロードして見比べて確認する。
🐣
フリーランスエンジニアです。
お仕事のご相談こちらまで
rockyshikoku@gmail.com
機械学習、ARアプリ（Web/iOS）を作っています。
機械学習／AR関連の情報を発信しています。