labelmeのjsonファイルを使ってセグメンテーション用データセットを作る方法

こんにちは、タナカです。

この記事では、labelmeのアノテーション結果を使ってセグメンテーション用データセットの作り方を紹介します。

アノテーションツールlabelmeの使い方という記事でアノテーションする方法を解説しましたが、今回は実際にアノテーション結果を使ってセマンティックセグメンテーション用のデータセットを実装していきます。

pythonを勉強していた当初は、labelmeでアノテーションしたものの、この結果を使ってセマンティックセグメンテーションを実装する方法が分かりませんでした。

特にjsonというこれまでなじみのなかった拡張子のファイルを扱うこともあり、少し毛嫌いするところもありましたね。

それでもlabelmeを使って、セマンティックセグメンテーションを実装するなら、jsonファイルの扱いに慣れる必要があります。

自分の整理も兼ねてできるだけわかりやすく解説できればと思います。

1. labelmeとは
2. セマンティックセグメンテーションとは
3. labelmeで作ったjsonファイルの中身を確認する
4. データセットクラスを定義する
まとめ

1. labelmeとは

labelmeとはセマンティックセグメンテーションを実装するときなどに使用するオープンソースのアノテーションツールになります。GUIで直感的に操作することが可能です。詳細についてはアノテーションツールlabelmeの使い方で紹介してます。

2. セマンティックセグメンテーションとは

セマンティックセグメンテーションとは、ピクセル単位で各物体を検出する機械学習アルゴリズムの一つです。

自動運転、医療用画像処理、工場での製品検査などに応用されています。

セマンティックセグメンテーションで学習を行う場合、下図のような画像を扱うことになります。左が学習データで、右が教師データになります。教師データが画像であることに注意してください。

セマンティックセグメンテーションを行うことで、最終的には画像中のどこに物体がいるか、その物体の面積はいくつかを算出できるようになります。

3. labelmeで作ったjsonファイルの中身を確認する

データセットを作る前にlabelmeで作ったjsonファイルの中身を確認していきます。jsonファイルの中身を確認するコードがこちらになります。

import json

with open('2021053020403400-194D89293F260C6893CF3FBF65B93019.json', 'r') as f:
    json_datas = json.load(f)

print(json_datas.keys())
# >> dict_keys(['version', 'flags', 'shapes', 'imagePath', 'imageData', 'imageHeight', 'imageWidth'])

print(json_datas['version']) # >> 4.5.13
print(json_datas['imagePath']) # >> '2021053020403400-194D89293F260C6893CF3FBF65B93019.json'
print(json_datas['imageData']) # >> 画像データ（文字列）
print(json_datas['imageHeight']) # >> 720
print(json_datas['imageWidth']) # >> 1280

print(json_datas['shapes'][0].keys()) # 0は0番目の要素のこと(画像内の物体の数だけ要素が存在)
# >> dict_keys(['label', 'points', 'group_id', 'shape_type', 'flags'])

print(json_datas['shapes'][0]['points'])
# >> アノテーションした座標

print(json_datas['shapes'][0]['label'])
# >> 自分でつけたラベル名

標準ライブラリであるjsonライブラリをインポートして、jsonファイルの中身をjson_datasという変数に入れています。

json_datasは辞書のような形式をしており、様々なkeyを保持しています。具体的なものでいうと画像データや画像の大きさなどになります。

json_datasの中にはさらに辞書が入っており、shapesをkeyにすることでlabelmeでアノテーションした座標データやラベル名などを取得することができます。

jsonデータの中身を把握しておけば、データセットを作る時に理解が深まると思います。

4. データセットクラスを定義する

ここからが本題のセマンティックセグメンテーションを実装するためのデータセットの作り方を解説していきます。早速コードはこちらになります。

# 標準ライブラリ
import json
import base64
import glob

# 外部ライブラリ
import numpy as np
from PIL import Image, ImageDraw
from labelme import utils
from torch.utils.data import Dataset, DataLoader
import albumentations as albu
from albumentations.pytorch import ToTensorV2

class MyDatasets(Dataset):
    def __init__(self, json_dir, phase, img_size):
        self.imgs, self.masks = self.createDatasets(json_dir)
        self.transform = self.albu_aug(phase, img_size)

    def __len__(self):
        return len(self.imgs)

    def __getitem__(self, index):
        img = self.imgs[index] / 255
        mask = self.masks[index]

        datasets = self.transform(image = img, mask = mask)
        img = datasets['image']
        mask = datasets['mask'].permute(2, 0, 1)
        img = img.type(torch.float32)
        mask = mask.type(torch.float)
        return img, mask

    # albumentationsによる水増し関数
    def albu_aug(self, phase, img_size):
        if phase == 'train':
            augmentation = albu.Compose([
                albu.LongestMaxSize(max_size = img_size, always_apply = True),
                albu.HorizontalFlip(p = 0.5),
                ToTensorV2()
            ])
        else:
            augmentation = albu.Compose([
                albu.LongestMaxSize(max_size = img_size, always_apply = True),
                ToTensorV2()
            ])
        return augmentation        

    def createDatasets(self, json_dir):
        # imgs
        imgs = []
        # masks
        masks = []

        # jsonデータの読み込み
        json_paths = glob.glob(json_dir + '/*.json')
        for json_path in json_paths:
            json_file = open(json_path)
            json_data = json.load(json_file)

            # imageDataをkeyにしてデータを取り出す
            img_b64 = json_data['imageData']

            # labelmeのutils関数を使ってbase64形式をPIL型に変換する
            img_data = base64.b64decode(img_b64)
            img_pil = utils.img_data_to_pil(img_data)
            img_arr = np.array(img_pil)

            # imgsに画像を追加
            imgs.append(img_arr)

            # マスク画像準備
            w, h = img_pil.size
            mask = Image.new('L', (w, h))

            # 物体の数で反復処理
            for i in range(len(json_data['shapes'])):
                # 座標情報
                points = json_data['shapes'][i]['points']

                # 物体を色塗り（クラス数は1を想定）
                draw = ImageDraw.Draw(mask)
                draw.polygon([tuple(point) for point in points], fill = 1, outline = 1)

            # mask array
            mask_arr = np.array(mask)

            # マスク画像生成
            obj = np.unique(mask_arr)[0:]
            labels = np.array(obj)
            mask = np.zeros((mask_arr.shape[0], mask_arr.shape[1], len(labels)), dtype = 'uint8')
            for i in range(len(labels)):
                temp = np.where((mask_arr[:, :] == labels[i]), 1, 0)
                mask[:, :, i] = temp
            masks.append(mask)

        return imgs, masks

少し補足をするとalbumentationsのToTensorV2は、array型をtensor型に変換する処理になります。このときに画像の画素値の最大である255で割る処理をしていないため、__getitem__のところで255で割る処理を追加しています。

1つ注意点として、pytorchのtransformsという関数にも、ToTensorという似たような処理があるのですが、これはtensor型と同時に255で割る処理をしているので、間違えないようにしたいところです。

なぜ、ToTensorV2を使用しているかというと、マスク画像は0, 1で表現されているので、マスク画像を255で割らないため（元画像とマスク画像で処理を分けるため）に、今回はToTensorV2を使用しています。