<麻雀牌検出器作成記①>麻雀牌教師データセットの生成【オブジェクト検出】

from PIL import Image

# 画像正方形化関数(黒余白)

def expand2square(pil_img, background_color=(0, 0, 0)):
    width, height = pil_img.size
    if width == height:
        return pil_img
    elif width > height:
        result = Image.new(pil_img.mode, (width, width), background_color)
        result.paste(pil_img, (0, (width - height) // 2))
        return result
    else:
        result = Image.new(pil_img.mode, (height, height), background_color)
        result.paste(pil_img, ((height - width) // 2, 0))
        return result

<annotation>
    <!-- 画像データの場所 -->
    <folder>xxx</folder>
    <filename>xxx.jpg</filename>

    <!-- 画像データの詳細(サイズ、次元) -->
    <size>
        <width>1000</width>
        <height>1000</height>
        <depth>3</depth>
    </size>

    <!-- ここから各オブジェクトの情報 -->
    <object>
        <name>circle_1</name> <!-- オブジェクトのラベル(例は一筒) -->
        <difficult>0</difficult>
        <bndbox>
            <xmin>100</xmin> <!-- 開始X座標(左) -->
            <ymin>100</ymin> <!-- 開始Y座標(上) -->
            <xmax>160</xmax> <!-- 終了X座標(右) -->
            <ymax>180</ymax> <!-- 終了Y座標(下) -->
        </bndbox>
    </object>

    <object>
        <name>character_9</name> <!-- 例は九萬 -->
        <difficult>0</difficult>
        <bndbox>
            <xmin>200</xmin>
            <ymin>200</ymin>
            <xmax>260</xmax>
            <ymax>280</ymax>
        </bndbox>
    </object>
</annotation>

import cv2
from IPython.display import Image, display
import random
import os
import json

# アスペクト比を固定したまま、imgを指定したwidthにリサイズする関数
def scale_to_width(img, width):
    h, w = img.shape[:2]
    height = round(h * (width / w)) # リサイズ後のheightを計算
    dst = cv2.resize(img, dsize=(width, height))

    return dst

# 配置したい麻雀牌のサイズ(リストに指定した中から、ランダムに1つ選ばれる)
sizes = [51, 54, 57, 60, 63]
# データセットJSONに記載するラベル(上から筒子、索子、萬子、字牌)
labels = ["circle_1", "circle_2", "circle_3", "circle_4", "circle_5", "circle_6", "circle_7", "circle_8", "circle_9", \
          "bamboo_1", "bamboo_2", "bamboo_3", "bamboo_4", "bamboo_5", "bamboo_6", "bamboo_7", "bamboo_8", "bamboo_9", \
          "chara_1", "chara_2", "chara_3", "chara_4", "chara_5", "chara_6", "chara_7", "chara_8", "chara_9", \
          "east", "south", "west", "north", "red", "green", "white"]

width_size = random.choice(sizes)

# 1枚単体で大きく映った麻雀牌画像のディレクトリパス
# このディレクトリの直下に、さらにそれぞれラベル(牌)ごとのディレクトリを配置し、
# その各ディレクトリの下にそのラベル(牌)の画像を配置する
dataset_dir_path = os.listdir("./dataset_org/")

json_dump_list = []

back_img_size = 512

# 1000枚のデータセットを作る
for i in range(1000):
    loop_i = 0
    # 麻雀牌画像を配置する背景画像のディレクトリパス(麻雀マットをイメージして、緑一面の画像とかでよい)
    # 512 × 512 の画像とする
    back_img = cv2.imread("./mat.jpg")
    image_infos = []

    # 1画像あたり、10枚の麻雀牌を配置する
    for _ in range(10):
        # 画像に配置するラベル(牌)をランダムに選ぶ
        dataset_path = random.choice(dataset_dir_path)
        images_path = f'./dataset_org/{dataset_path}'
        images = os.listdir(images_path)
        image_path =f'{images_path}/{random.choice(images)}'
    
        fore_img = cv2.imread(image_path)

        # 画像のサイズをアスペクトを保持したままリサイズする
        fore_img = scale_to_width(fore_img, width_size)

        # 50%の確立で、180°回転させる
        if random.random() > 0.5:
            fore_img = cv2.rotate(fore_img, cv2.ROTATE_180)
        
        # 10%の確立で、90°回転させる
        if random.random() > 0.9:
            fore_img = cv2.rotate(fore_img, cv2.ROTATE_90_CLOCKWISE)

        # 配置する麻雀牌のサイズ(リサイズ後)
        fore_h, fore_w, fore_c = fore_img.shape

        # 配置場所開始候補(x, y)をランダムに決める
        rand_x = random.randint(0, back_img_size-1-fore_w)
        rand_y = random.randint(0, back_img_size-1-fore_h)

        # 50%の確立で前回の画像と隣に配置する(512の背景画像をはみ出さなければ)
        if len(image_infos) > 0 and random.random() >= 0.5 and\
            image_infos[-1]["end_x"] + fore_w +1 <= back_img_size-1 and image_infos[-1]["start_y"] + fore_h <= back_img_size-1:
            # 開始候補を隣の位置に更新
            rand_x = image_infos[-1]["end_x"] +1
            rand_y = image_infos[-1]["start_y"]    

        # 配置失敗時(画像をはみ出すとき、既に配置した画像と重複するとき)のために、
        # 100回までの間で、成功するまで試行する
        while loop_i < 100:
            ng_flag = False # 配置失敗時のフラグ
            # 既に配置した麻雀牌と重複するかどうかをチェックする
            for info in image_infos:
                if (rand_x >= info["start_x"] and rand_x <= info["end_x"]) or\
                    (rand_x + fore_w >= info["start_x"] and rand_x + fore_w <= info["end_x"]):
                        if (rand_y >= info["start_y"] and rand_y <= info["end_y"]) or\
                            (rand_y + fore_h >= info["start_y"] and rand_y + fore_h <= info["end_y"]) or\
                            (info["start_y"] >= rand_y and info["start_y"] <= rand_y + fore_h) or\
                            (info["end_y"] >= rand_y and info["end_y"] <= rand_y + fore_h):
                            ng_flag = True # 失敗時はNG

            loop_i += 1        

            # NGフラグがONのときは、配置場所開始候補をランダムに再決定する
            if ng_flag:
                rand_x = random.randint(0, back_img_size-1-fore_w)
                rand_y = random.randint(0, back_img_size-1-fore_h)
                continue

            # NGフラグがOFF(配置成功)のときは、1枚の牌の配置情報を追加し、
            # 実際に画像に牌を上塗りし、100回チャンレジから抜ける
            else:
                image_infos.append({
                    "num": dataset_path,
                    "start_x": rand_x,
                    "end_x": rand_x + fore_w, 
                    "start_y": rand_y,
                    "end_y": rand_y + fore_h
                })
                back_img[rand_y:rand_y+fore_h, rand_x:rand_x+fore_w] = fore_img
                break
       
    # 10枚の麻雀牌を配置し終えたら、1枚の画像として保存する
    cv2.imwrite(f'./dataset/{i}.jpg', back_img)
    
    # JSONに情報を追加する
    # 一番上に画像ファイル名
    json_dump_list.append({
        "image": f'{i}.jpg',
        "annotations": []
    })
    # JSONのannotationsに配置した麻雀牌の座標情報を追加していく
    for info in image_infos:
        json_dump_list[-1]["annotations"].append(
            {
                "label": labels[int(info["num"])-1],
                "coordinates": {
                    "x": info["start_x"] + (info["end_x"] - info["start_x"])/2,
                    "y": info["start_y"] + (info["end_y"] - info["start_y"])/2 ,
                    "width": info["end_x"]-info["start_x"],
                    "height": info["end_y"]-info["start_y"]
                }
            }
        )

# JSONファイルを保存する
with open('./dataset.json', 'w') as f:
    json.dump(json_dump_list, f)