Pythonあれこれ（Vol.92掲載）
著者：飯尾 淳
本連載では「Pythonを昔から使っているものの、それほど使いこなしてはいない」という筆者が、いろいろな日常業務をPythonで処理することで、立派な「蛇使い」に育つことを目指します。その過程を温かく見守ってください。皆さんと共に勉強していきましょう。第22回では、機械学習ライブラリ「PyTorch」を使って、画像内の、歩行者が写っている領域を自動認識するPythonプログラムを作成します。
シェルスクリプトマガジン Vol.92は以下のリンク先でご購入できます。
図3 47番のデータを確認するPythonコード
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 |
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コード
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
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コード
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 |
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コード
|
1 2 3 4 5 6 7 8 9 |
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コード
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
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コード
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 |
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コード
|
1 2 3 4 5 6 7 8 9 10 |
# 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) |