Pytorchには学習済みモデルが実装されており関数を呼び出すだけですぐに使えるようになっています。
「VGG16」や「Resnet18」などは画像データをTensorに変換したデータを入力すれば出力を得ることができますが、 「Faster R-CNN」のような複雑なモデルは入力画像とDict形式のTargetデータが必要になり、所定の形式でなければ エラーになってしまいます。
今回はLabelmeで出力したデータセットから、Pytorchに実装されているFaster R-CNNで使えるデータセットの作成方法についてまとめました。
1. カスタムデータセットの作成
1.1. カスタムデータセットのフォーマット
まずはカスタムデータセットの作成についてです。 データセットを自作する場合は「torch.utils.data.Dataset」を継承したクラスを作成し、以下の関数を追加します。
- __init__ :オブジェクト作成時に呼び出される
- __len__:len()関数を実行したときに呼び出される(データセットのサイズを取得)
- __getitem__:インデックスキーを指定して要素を取得するときに呼び出される（指定のインデックス番号のデータセットを取得）
import torch
class Custom_Dataset(torch.utils.data.Dataset):
def __init__(self, transforms):
self.imgs = []
self.targets = []
def __len__(self):
return len(self.imgs)
def __getitem__(self, idx):
if self.transforms is not None:
img, target = self.transforms(img, target)
return img, target1.2. Jsonファイルを読み込みデータセットを作成
次にLabelmeから出力されたJsonファイルを読み込みデータセットを作成する関数を実装します。
今回はNumpy型の画像データとDict型のtargetsデータを作成しました。
#インポート追加
import json
from PIL import Image
import base64
import glob
import io
# base64形式をPIL型に変換
def img_data_to_pil(self,img_data):
f = io.BytesIO()
f.write(img_data)
img_pil = Image.open(f)
return img_pil
def CreateDataset(self, json_dir):
# 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']
#base64形式をPIL型に変換する
img_data = base64.b64decode(img_b64)
img_pil = self.img_data_to_pil(img_data)
# imgsに画像を追加
self.imgs.append(img_pil)
# targets作成
num_objs = len(json_data['shapes']) # 物体の個数
boxes = []
# バウンティングボックスから左上、左下、右上、右下の座標を取得し、
# boxesリストへ追加
for i in range(num_objs):
box = json_data['shapes'][i]['points']
x_list = []
y_list = []
for i in range(len(box)):
x,y=box[i]
x_list.append(x)
y_list.append(y)
xmin = min(x_list)
xmax = max(x_list)
ymin = min(y_list)
ymax = max(y_list)
boxes.append([xmin, ymin, xmax, ymax])
# boxesリストをtorch.Tensorに変換
boxes = torch.as_tensor(boxes, dtype=torch.float32)
# torch.Tensorにラベルを設定
# 物体の個数(背景を0とする)
labels = torch.ones((num_objs,), dtype=torch.int64)
#targetsへboxesリストとlabelsリストを設定
target = {}
target["boxes"] = boxes
target["labels"] = labels
self.targets.append(target)1.3. データセットクラス全体
データセットクラス全体は以下のようになります。
※ targetはDeepCopyでないと前処理が初期化されずに残ってしまうようです。
※2025.10.11追記
import copyが抜けいていました。
import torch
import json
from PIL import Image
import base64
import glob
import io
import copy
class Custom_Dataset(torch.utils.data.Dataset):
def __init__(self, root, transforms):
self.imgs = []
self.targets = []
self.transforms = transforms
self.CreateDataset(root)
def __len__(self):
return len(self.imgs)
def __getitem__(self, idx):
img = self.imgs[idx]
target = copy.deepcopy(self.targets[idx])
if self.transforms is not None:
img, target = self.transforms(img, target)
return img, target
# base64形式をPIL型に変換
def img_data_to_pil(self,img_data):
f = io.BytesIO()
f.write(img_data)
img_pil = Image.open(f)
return img_pil
def CreateDataset(self, json_dir):
# 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']
#base64形式をPIL型に変換する
img_data = base64.b64decode(img_b64)
img_pil = self.img_data_to_pil(img_data)
# imgsに画像を追加
self.imgs.append(img_pil)
# targets作成
num_objs = len(json_data['shapes']) # 物体の個数
boxes = []
# バウンティングボックスから左上、左下、右上、右下の座標を取得し、
# boxesリストへ追加
for i in range(num_objs):
box = json_data['shapes'][i]['points']
x_list = []
y_list = []
for i in range(len(box)):
x,y=box[i]
x_list.append(x)
y_list.append(y)
xmin = min(x_list)
xmax = max(x_list)
ymin = min(y_list)
ymax = max(y_list)
boxes.append([xmin, ymin, xmax, ymax])
# boxesリストをtorch.Tensorに変換
boxes = torch.as_tensor(boxes, dtype=torch.float32)
# torch.Tensorにラベルを設定
# 物体の個数(背景を0とする)
labels = torch.ones((num_objs,), dtype=torch.int64)
#targetsへboxesリストとlabelsリストを設定
target = {}
target["boxes"] = boxes
target["labels"] = labels
self.targets.append(target)2. 前処理(transforms)の実装
次に前処理(transforms)の実装を行います。
データセットをモデルに入力する前にリサイズや水増し(データのクロップや反転など)を行う場合があり、これらの処理をtransformsへ追加しておくことで データセットからデータを取り出すときに前処理された状態で取得することができます。
またPIL形式の画像データはTensor形式に変換する必要があるのですが、この変換処理もtransformsに追加しておきます。
#インポート追加
from torchvision.transforms import functional as F
class ToTensor(object):
def __call__(self, image, target):
image = F.to_tensor(image)
return image, target
class Compose(object):
def __init__(self, transforms):
self.transforms = transforms
def __call__(self, image, target):
for t in self.transforms:
image, target = t(image, target)
return image, target
def get_transform(train):
transforms = []
# PIL imageをPyTorch Tensorに変換
transforms.append(ToTensor())
if train:
#学習時のみの処理を追加
pass
return Compose(transforms)3. データセットの確認
次に、データセットからデータを取り出してみます。
今回はこちらの画像のロゴ部分にLabelmeでアノテーションを行いました。
Labelmeから出力されたデータはこちら。
{
"version": "5.5.0",
"flags": {},
"shapes": [
{
"label": "logo",
"points": [
[
44.451219512195124,
89.28455284552845
],
[
416.4024390243902,
269.7723577235772
]
],
"group_id": 1,
"description": "",
"shape_type": "rectangle",
"flags": {},
"mask": null
}
],
"imagePath": "shimons_labo1.jpg",
...データセットクラスにJsonファイルが入っているフォルダとtransformsを指定し、出力結果を表示します。
import torchvision
from PIL import ImageDraw
#1章で作成したデータセットをインポート
from CustomDataset import Custom_Dataset
#2章で作成transformsをインポート
from transforms import get_transform
dataset = Custom_Dataset('dataset',get_transform(train=True))
img, target = dataset[0]
# Tensor -> PILへ変換
img = torchvision.transforms.functional.to_pil_image(img)
# バウンティングボックス描画
for i in range(len(target['boxes'])):
draw = ImageDraw.Draw(img)
x = target['boxes'][i][0]
y = target['boxes'][i][1]
w = target['boxes'][i][2]
h = target['boxes'][i][3]
draw.rectangle((x,y,w,h), outline=(255, 0, 0), width=3)
print(target)
img.show()出力結果は以下の通り。
{'boxes': tensor([[ 44.4512, 89.2846, 416.4024, 269.7724]]), 'labels': tensor([1])}
4. Faster R-CNNで出力を確認
Pytorchに実装されているFaster R-CNNへ入力して出力を確認してみます。
Pytorchのチュートリアルを元にFaster R-CNNのモデルを取得する関数を実装します。
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
def get_instance_model(num_classes):
# COCOデータセットで訓練した、訓練済みモデルをロード
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False)
# 既存の分類器を、ユーザーが定義したnum_classesを持つ新しい分類器に置き換えます
num_classes = num_classes # 人を示すクラス＋背景クラスで2個
# 分類器にインプットする特徴量の数を取得
in_features = model.roi_heads.box_predictor.cls_score.in_features
# 事前訓練済みのヘッドを新しいものと置き換える
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
return modelデータセットからデータを取り出し、Faster R-CNNへ入力して動作を確認します。
※ビット深度が24じゃないとエラーが出てしまいます。
#1章で作成したデータセットをインポート
from CustomDataset import Custom_Dataset
#2章で作成transformsをインポート
from transforms import get_transform
#2025.10.11追記
from model import get_instance_model
#データセット読み込み
dataset = Custom_Dataset('dataset',get_transform(train=True))
img,target = dataset[0]
# モデルへ入力するためにListへ追加
images = []
targets = []
images.append(img)
targets.append(target)
print(images)
print(targets)
#モデルの定義
model = get_instance_model(num_classes=2)
#計算
output = model(images, targets)
#出力表示
for name,data in output.items():
print(name, data.item())出力結果はこちら。
loss_classifier 0.7495388984680176
loss_box_reg 5.054557550465688e-05
loss_objectness 0.6973650455474854
loss_rpn_box_reg 0.0029146710876375437自作のデータセットでFaster R-CNNの出力が得られることが確認できました。
今回は以上です。
5. 参考文献
・labelmeのjsonファイルを使ってセグメンテーション用データセットを作る方法
https://tanalib.com/labelme-to-dataset/
・「Torchvisionを利用した物体検出のファインチューニング手法」
https://colab.research.google.com/github/YutaroOgawa/pytorch_tutorials_jp/blob/main/notebook/2_Image_Video/2_2_torchvision_finetuning_instance_segmentation_jp.ipynb
コメント