Pythonあれこれ（Vol.92掲載）
著者：飯尾 淳
本連載では「Pythonを昔から使っているものの、それほど使いこなしてはいない」という筆者が、いろいろな日常業務をPythonで処理することで、立派な「蛇使い」に育つことを目指します。その過程を温かく見守ってください。皆さんと共に勉強していきましょう。第22回では、機械学習ライブラリ「PyTorch」を使って、画像内の、歩行者が写っている領域を自動認識するPythonプログラムを作成します。
シェルスクリプトマガジン Vol.92は以下のリンク先でご購入できます。
図3 47番のデータを確認するPythonコード
import matplotlib.pyplot as plt
from torchvision.io import read_image
image =read_image(
"data/PennFudanPed/PNGImages/FudanPed00047.png")
mask = read_image(
"data/PennFudanPed/PedMasks/FudanPed00047_mask.png")
plt.figure(figsize=(16, 8))
plt.subplot(121)
plt.title("Image")
plt.imshow(image.permute(1, 2, 0))
plt.subplot(122)
plt.title("Mask")
plt.imshow(mask.permute(1, 2, 0))図5 PennFudanDataSetクラスを定義するPythonコード
import os
import torch
from torchvision.io import read_image
from torchvision.ops.boxes import masks_to_boxes
from torchvision import tv_tensors
from torchvision.transforms.v2 import functional as F
class PennFudanDataset(torch.utils.data.Dataset):
def __init__(self, root, transforms):
self.root = root
self.transforms = transforms
# イメージデータとマスクデータをロードし、ソートして並べておく
self.imgs = \
list(sorted(os.listdir(os.path.join(root, "PNGImages"))))
self.masks = \
list(sorted(os.listdir(os.path.join(root, "PedMasks"))))
def __getitem__(self, idx):
# イメージとマスクのロード
img_path = \
os.path.join(self.root, "PNGImages", self.imgs[idx])
mask_path = \
os.path.join(self.root, "PedMasks", self.masks[idx])
img = read_image(img_path)
mask = read_image(mask_path)
# インスタンスは色別にエンコードされていて……
obj_ids = torch.unique(mask)
# 最初のIDは背景色なので削除
obj_ids = obj_ids[1:]
num_objs = len(obj_ids)
# 色別にエンコードされているマスクを2値マスクに分ける
masks = \
(mask == obj_ids[:, None, None]).to(dtype=torch.uint8)
# マスクデータに対してバウンディングボックスを求める
boxes = masks_to_boxes(masks)
labels = torch.ones((num_objs,), dtype=torch.int64)
image_id = idx
area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
# 群衆でないと仮定
iscrowd = torch.zeros((num_objs,), dtype=torch.int64)
# tv_tensorsイメージに変換する
img = tv_tensors.Image(img)
target = {}
target["boxes"] = \
tv_tensors.BoundingBoxes(boxes,
format="XYXY", canvas_size=F.get_size(img))
target["masks"] = tv_tensors.Mask(masks)
target["labels"] = labels
target["image_id"] = image_id
target["area"] = area
target["iscrowd"] = iscrowd
if self.transforms is not None:
img, target = self.transforms(img, target)
return img, target
def __len__(self):
return len(self.imgs)図6 get_model_instance_segmentation()関数を定義するPythonコード
import torchvision
from torchvision.models.detection.faster_rcnn \
import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn \
import MaskRCNNPredictor
def get_model_instance_segmentation(num_classes):
# COCOで事前学習済みのモデルデータをロードする
model = torchvision.models.detection.maskrcnn_resnet50_fpn(
weights="DEFAULT")
# 入力特徴量
in_features = \
model.roi_heads.box_predictor.cls_score.in_features
# num_classesで指定するクラスの分類器をセット（今回は2クラス）
model.roi_heads.box_predictor = \
FastRCNNPredictor(in_features, num_classes)
# マスク判別器も同様に設定
in_features_mask = \
model.roi_heads.mask_predictor.conv5_mask.in_channels
hidden_layer = 256
model.roi_heads.mask_predictor = MaskRCNNPredictor(
in_features_mask,
hidden_layer,
num_classes
)
return model図7 get_transform()関数を定義するPythonコード
from torchvision.transforms import v2 as T
def get_transform(train):
transforms = []
if train:
transforms.append(T.RandomHorizontalFlip(0.5))
transforms.append(T.ToDtype(torch.float, scale=True))
transforms.append(T.ToPureTensor())
return T.Compose(transforms)図8 追加学習の準備をするためのPythonコード
import utils
from engine import train_one_epoch, evaluate
device = torch.device('cuda') \
if torch.cuda.is_available() else torch.device('cpu')
# 背景と人物の2クラス判別器を作成する
num_classes = 2
# データセットはすでにロード済みのものを用いる
dataset = \
PennFudanDataset('data/PennFudanPed',
get_transform(train=True))
dataset_test = \
PennFudanDataset('data/PennFudanPed',
get_transform(train=False))
# データセットを学習用とテスト用に分ける
indices = torch.randperm(len(dataset)).tolist()
dataset = torch.utils.data.Subset(dataset, indices[:-50])
dataset_test = torch.utils.data.Subset(
dataset_test, indices[-50:])
# 学習と評価用のデータローダを定義
data_loader = torch.utils.data.DataLoader(
dataset,
batch_size=2,
shuffle=True,
num_workers=4,
collate_fn=utils.collate_fn
)
data_loader_test = torch.utils.data.DataLoader(
dataset_test,
batch_size=1,
shuffle=False,
num_workers=4,
collate_fn=utils.collate_fn
)
# 先ほど定義したヘルパーファンクションを用いてモデルを用意
model = get_model_instance_segmentation(num_classes)
# モデルをデバイスに結び付ける
model.to(device)
# 最適化器を作成
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(
params,
lr=0.005,
momentum=0.9,
weight_decay=0.0005
)
# 学習レートのスケジューラを設定
lr_scheduler = torch.optim.lr_scheduler.StepLR(
optimizer,
step_size=3,
gamma=0.1
)図9 物体認識と領域セグメンテーションを実施するPythonコード
import matplotlib.pyplot as plt
from torchvision.utils \
import draw_bounding_boxes, draw_segmentation_masks
image = \
read_image("data/PennFudanPed/PNGImages/FudanPed00047.png")
eval_transform = get_transform(train=False)
model.eval()
with torch.no_grad():
x = eval_transform(image)
# RGBA -> RGBにコンバートしてデバイスにひも付ける
x = x[:3, ...].to(device)
predictions = model([x, ])
pred = predictions[0]
image = (255.0 * (image - image.min()) /
(image.max() - image.min())).to(torch.uint8)
image = image[:3, ...]
pred_labels = [f"pedestrian: {score:.3f}" for label, \
score in zip(pred["labels"], pred["scores"])]
pred_boxes = pred["boxes"].long()
output_image = draw_bounding_boxes(image,
pred_boxes, pred_labels, colors="red")
masks = (pred["masks"] > 0.7).squeeze(1)
output_image = draw_segmentation_masks(output_image,
masks, alpha=0.5, colors="blue")
plt.figure(figsize=(12, 12))
plt.imshow(output_image.permute(1, 2, 0))図11 Penn-Fudanデータセットを用いて追加学習するPythonコード
# 2エポックだけ学習させてみる
num_epochs = 2
for epoch in range(num_epochs):
# 10回ごとに表示させながら1エポックの学習を実行
train_one_epoch(model, optimizer, data_loader,
device, epoch, print_freq=10)
# 学習レートをアップデート
lr_scheduler.step()
# テストデータで評価
evaluate(model, data_loader_test, device=device)