PyTorch入門 第9回：画像分類、セグメンテーション、物体検出をまとめて理解する

第8回：モデルを扱える状態にする
  -> state_dictを保存する
  -> checkpointから再開する
  -> 事前学習済みモデルを読み込む
  -> deviceへ移して安全に実行する

第9回：画像タスクの出力を理解する
  -> 分類logitsをクラス名へ変換する
  -> セグメンテーションlogitsをmaskへ変換する
  -> 検出結果のbox、label、scoreを読む
  -> IoU、NMS、mAPなどの評価と後処理を理解する

単一画像: (C, H, W)
batch  : (N, C, H, W)

logits shape = (N, num_classes)

class_to_idx = {"ants": 0, "bees": 1}  # 訓練時に決めたクラス名 -> クラスIDの対応表
idx_to_class = {value: key for key, value in class_to_idx.items()}  # 逆引き用の対応表（クラスID -> クラス名）

print(idx_to_class[1])
# bees

import torch


categories = ["cat", "dog", "bird"]
logits = torch.tensor(
    [[1.0, 3.0, 0.2],
     [2.5, 0.1, 0.0]]
)

predicted = logits.argmax(dim=1)  # クラス次元(dim=1)で最大のindexを取る -> 予測クラスID
probabilities = logits.softmax(dim=1)  # クラス次元を確率らしい値に正規化する（各サンプルの合計が1になる）
confidence = probabilities.max(dim=1).values  # サンプルごとに最大確率だけを取り出す
names = [categories[index] for index in predicted.tolist()]  # クラスIDを文字列名に変換する

print("predicted:", predicted.tolist())
# predicted: [1, 0]
print("names:", names)
# names: ['dog', 'cat']
print("confidence:", [round(value.item(), 4) for value in confidence])
# confidence: [0.836, 0.8527]

単一ラベル分類: softmax -> argmax
マルチラベル分類: sigmoid -> classごとのthreshold

画像を読み込む
 -> RGB/BGRを確認する
 -> Resize / Crop
 -> Tensorへ変換し、値域を整える
 -> Normalize
 -> batch次元を追加する
 -> deviceへ移す
 -> model.eval()
 -> torch.inference_mode()
 -> logits
 -> argmax / top-k
 -> idx_to_class

通常の深いblock:
  output = H(x)

残差block:
  F(x) = H(x) - x
  output = F(x) + x

x ------------------------------+
 |                               |
 +-> Conv -> Norm -> ReLU -> Conv+
                                 |
                         Add -> ReLU

x -> 1x1 Conv / Norm -> shortcut

from PIL import Image
import torch
from torchvision.models import resnet18, ResNet18_Weights


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

weights = ResNet18_Weights.DEFAULT  # 推奨される既定の事前学習済み重みを選ぶ
model = resnet18(weights=weights).eval().to(device)
preprocess = weights.transforms()  # その重みが学習時に使った前処理（Resize/Normalizeなど）を取得する
categories = weights.meta["categories"]  # 重みに紐づくImageNet 1000クラスの名前一覧

image = Image.open("your_image.jpg").convert("RGB")  # 自分の画像ファイルに置き換える
input_tensor = preprocess(image).unsqueeze(0).to(device)

with torch.inference_mode():
    logits = model(input_tensor)
    probabilities = logits.softmax(dim=1)
    top_prob, top_index = probabilities.topk(5, dim=1)  # 確率が高い上位5件を取り出す

for probability, index in zip(top_prob[0], top_index[0]):
    print(categories[index], float(probability))  # クラス名と確率をスコアの高い順に表示する

from pathlib import Path
import time
import torch
from torch import nn
from PIL import Image
from torchvision import transforms
from torchvision.models import resnet18


classes = ["ants", "bees"]
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

inference_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225],
    ),
])


def build_resnet18_for_two_classes(checkpoint_path):
    model = resnet18(weights=None)  # checkpoint側の重みを読み込むので、ここでは構造だけ作る
    model.fc = nn.Linear(model.fc.in_features, len(classes))  # 1000クラス用headを2クラス用に差し替える

    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
    state_dict = checkpoint.get("model_state_dict", checkpoint)  # 辞書に包まれたcheckpointと、state_dict直接の両方に対応する
    model.load_state_dict(state_dict)
    return model.to(device).eval()  # 推論用にdeviceへ移し、eval()でDropout/BatchNormを推論モードにする


def predict_image(image_path, checkpoint_path):
    image = Image.open(image_path).convert("RGB")
    image_tensor = inference_transform(image).unsqueeze(0).to(device)  # (C, H, W) -> (1, C, H, W)
    model = build_resnet18_for_two_classes(checkpoint_path)

    with torch.inference_mode():
        start = time.perf_counter()
        logits = model(image_tensor)
        elapsed = time.perf_counter() - start

    predicted_index = logits.argmax(dim=1).item()  # バッチ内に1枚だけなので.item()でPythonの数値に変換する
    return classes[predicted_index], elapsed


predicted_name, elapsed = predict_image(
    Path("your_image.jpg"),
    Path("checkpoint_14_epoch.pkl"),
)
print(f"predict: {predicted_name} elapsed: {elapsed:.3f}s")

import torch
from torchvision.models import resnet18


model = resnet18(weights=None).eval()
example = torch.randn(2, 3, 224, 224)

with torch.inference_mode():
    logits = model(example)

print("logits shape:", tuple(logits.shape))
# logits shape: (2, 1000)

import torch
from torch import nn


logits = torch.tensor(
    [[1.0, 3.0, 0.2],
     [2.5, 0.1, 0.0]]
)
targets = torch.tensor([1, 0])  # 各サンプルの正解クラスID（shape (N,) のtorch.long）

criterion = nn.CrossEntropyLoss()
loss = criterion(logits, targets)

print("loss:", round(loss.item(), 4))
# loss: 0.1693

import torch


logits = torch.tensor(
    [[0.1, 2.0, 1.0],
     [3.0, 1.0, 0.0]]
)
targets = torch.tensor([1, 1])

top1 = (logits.argmax(dim=1) == targets).float().mean()  # 予測の1位が正解と一致する割合
top2_indices = logits.topk(2, dim=1).indices  # 各サンプルの上位2クラスのindexを取る
top2 = top2_indices.eq(targets.unsqueeze(1)).any(dim=1).float().mean()  # 正解が上位2位以内に入っている割合

print("top1:", float(top1), "top2:", float(top2))
# top1: 0.5 top2: 1.0

input: (N, 3, H, W)

logits: (N, C, H_out, W_out)

import torch


logits = torch.tensor(
    [[[[0.1, 2.0],
       [0.2, 0.3]],
      [[1.2, 0.5],
       [0.1, 1.5]],
      [[0.0, 0.1],
       [2.2, 0.4]]]]
)

pred_mask = logits.argmax(dim=1)  # (N, C, H, W) -> (N, H, W)

print("mask shape:", tuple(pred_mask.shape))
# mask shape: (1, 2, 2)
print(pred_mask)
# tensor([[[1, 0],
#          [2, 1]]])

target shape: (N, H, W)
target dtype: torch.long
target value: 0..C-1 または ignore_index

import torch
from torch import nn


logits = torch.randn(2, 3, 4, 4)  # 2枚、3クラス、4x4のピクセルlogits
target = torch.randint(0, 3, (2, 4, 4), dtype=torch.long)

criterion = nn.CrossEntropyLoss()
loss = criterion(logits, target)

print("loss is scalar:", loss.ndim == 0)
# loss is scalar: True

IoU_c = TP_c / (TP_c + FP_c + FN_c)
Dice_c = 2 * TP_c / (2 * TP_c + FP_c + FN_c)

import torch


pred = torch.tensor(
    [[0, 1, 1],
     [0, 1, 2]]
)
target = torch.tensor(
    [[0, 1, 0],
     [0, 2, 2]]
)

ious = []
for class_id in range(3):
    pred_c = pred == class_id  # このクラスだと予測した位置のbool mask
    target_c = target == class_id  # このクラスが正解の位置のbool mask
    intersection = (pred_c & target_c).sum()  # 両方でTrueの画素数 = TP
    union = (pred_c | target_c).sum()  # どちらかでTrueの画素数
    ious.append((intersection / union).item())

print([round(value, 4) for value in ious])
# [0.6667, 0.3333, 0.5]

Encoderで空間サイズを下げながら特徴を取る
 -> coarse score mapを作る
 -> upsamplingで空間サイズを戻す
 -> skip fusionで浅い特徴と深い特徴を混ぜる
 -> dense pixel prediction

Input
 -> Encoder 1 ----skip----------------> Decoder 1
 -> Encoder 2 ------skip-----------> Decoder 2
 -> Bottleneck
 -> Upsample + concat encoder feature
 -> segmentation logits

Backbone -> atrous features -> ASPP
                             |
low-level feature -> projection -> concat -> decoder -> logits

from PIL import Image
import torch
from torchvision.models.segmentation import (
    deeplabv3_resnet50,
    DeepLabV3_ResNet50_Weights,
)


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

weights = DeepLabV3_ResNet50_Weights.DEFAULT
model = deeplabv3_resnet50(weights=weights).eval().to(device)
preprocess = weights.transforms()

image = Image.open("your_image.jpg").convert("RGB")  # 自分の画像ファイルに置き換える
x = preprocess(image).unsqueeze(0).to(device)

with torch.inference_mode():
    output = model(x)
    logits = output["out"]
    mask = logits.argmax(dim=1)

print("logits shape:", tuple(logits.shape))
print("mask shape:", tuple(mask.shape))

import torch


def voc_palette(num_classes=21):
    base = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
    class_ids = torch.arange(num_classes)[:, None]
    return (class_ids * base % 255).to(torch.uint8)


mask = torch.tensor(
    [[0, 15],
     [8, 2]],
    dtype=torch.long,
)
palette = voc_palette()
rgb_mask = palette[mask]  # (H, W) のクラスIDを (H, W, 3) の色へ変換する

print("rgb mask:", tuple(rgb_mask.shape), rgb_mask.dtype)
# rgb mask: (2, 2, 3) torch.uint8

I = alpha * F + (1 - alpha) * B

import torch


alpha = torch.tensor([0.0, 0.5, 1.0]).view(3, 1)
foreground = torch.ones(3, 3)
background = torch.zeros(3, 3)
composite = alpha * foreground + (1 - alpha) * background

print(composite)
# tensor([[0.0000, 0.0000, 0.0000],
#         [0.5000, 0.5000, 0.5000],
#         [1.0000, 1.0000, 1.0000]])

import torch


def compute_dice(pred_mask, true_mask, eps=1e-7):
    pred_mask = pred_mask.bool()
    true_mask = true_mask.bool()
    intersection = (pred_mask & true_mask).sum()  # 両方でTrueの画素数
    return (2 * intersection + eps) / (pred_mask.sum() + true_mask.sum() + eps)  # epsはゼロ除算を防ぐための微小値


pred_mask = torch.tensor(
    [[1, 0, 1],
     [0, 1, 0]]
)
true_mask = torch.tensor(
    [[1, 0, 0],
     [0, 1, 1]]
)

dice = compute_dice(pred_mask, true_mask)
print("dice:", round(float(dice), 4))
# dice: 0.6667

import torch
from torch import nn
from torch.utils.data import DataLoader


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
mask_threshold = 0.5

train_dataset = PortraitDataset("data/PortraitDataset/train")
valid_dataset = PortraitDataset("data/PortraitDataset/valid")
train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True, drop_last=True)
valid_loader = DataLoader(valid_dataset, batch_size=1, shuffle=False)

model = UNet(in_channels=3, out_channels=1, init_features=32).to(device)  # 前景/背景の2値なので出力チャンネルは1つ
criterion = nn.BCEWithLogitsLoss()  # 1チャンネルlogits + 0/1ラベルの組み合わせに合う損失
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=150, gamma=0.1)

for epoch in range(num_epochs):
    model.train()
    for images, masks in train_loader:
        images = images.to(device)
        masks = masks.to(device).float()

        logits = model(images)
        loss = criterion(logits, masks)

        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        optimizer.step()

        pred_masks = torch.sigmoid(logits).ge(mask_threshold)  # logits -> 確率 -> しきい値で二値maskへ
        dice = compute_dice(pred_masks.cpu(), masks.cpu().ge(mask_threshold))  # 直近バッチのDiceを確認する例

    scheduler.step()  # optimizer.step()の後、エポック単位でschedulerを進める

from collections import OrderedDict
from pathlib import Path
import torch
from PIL import Image
from torchvision import transforms


def load_unet(checkpoint_path):
    model = UNet(in_channels=3, out_channels=1, init_features=32)
    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)

    cleaned_state = OrderedDict(
        (key.removeprefix("module."), value)
        for key, value in checkpoint["model_state_dict"].items()
    )
    model.load_state_dict(cleaned_state)
    return model.to(device).eval()


preprocess = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
])

model = load_unet(Path("checkpoint_399_epoch.pkl"))
image = Image.open("your_portrait.png").convert("RGB")
image_tensor = preprocess(image).unsqueeze(0).to(device)

with torch.inference_mode():
    logits = model(image_tensor)
    mask = torch.sigmoid(logits).ge(0.5)

{
    "boxes":  FloatTensor[num_objects, 4],
    "labels": Int64Tensor[num_objects],
    "scores": Tensor[num_objects],
}

[x1, y1, x2, y2]
0 <= x1 < x2 <= W
0 <= y1 < y2 <= H

import torch


boxes = torch.tensor(
    [[10.0, 20.0, 50.0, 80.0],
     [15.0, 30.0, 45.0, 70.0]]
)
widths = boxes[:, 2] - boxes[:, 0]
heights = boxes[:, 3] - boxes[:, 1]
areas = widths * heights

print("areas:", areas.tolist())
# areas: [2400.0, 1200.0]

from PIL import Image
import torch
from torchvision.models.detection import (
    fasterrcnn_resnet50_fpn_v2,
    FasterRCNN_ResNet50_FPN_V2_Weights,
)


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
weights = FasterRCNN_ResNet50_FPN_V2_Weights.DEFAULT
model = fasterrcnn_resnet50_fpn_v2(weights=weights).eval().to(device)
preprocess = weights.transforms()
categories = weights.meta["categories"]

image = Image.open("your_detection_image.jpg").convert("RGB")
image_tensor = preprocess(image).to(device)

with torch.inference_mode():
    prediction = model([image_tensor])[0]

keep = prediction["scores"] >= 0.5
for box, label, score in zip(
    prediction["boxes"][keep],
    prediction["labels"][keep],
    prediction["scores"][keep],
):
    class_name = categories[int(label)]
    print(class_name, round(float(score), 3), [round(float(value), 1) for value in box])

IoU(A, B) = area(A ∩ B) / area(A ∪ B)

scoreが高い順に並べる
 -> 最高scoreのboxを残す
 -> そのboxとIoUが高すぎるboxを消す
 -> 残りで繰り返す

import torch
from torchvision.ops import box_iou, nms


boxes = torch.tensor(
    [[0.0, 0.0, 10.0, 10.0],
     [1.0, 1.0, 11.0, 11.0],
     [20.0, 20.0, 30.0, 30.0]]
)
scores = torch.tensor([0.9, 0.8, 0.7])

ious = box_iou(boxes[:1], boxes)
keep = nms(boxes, scores, iou_threshold=0.5)

print(torch.round(ious, decimals=4))
# tensor([[1.0000, 0.6807, 0.0000]])
print("keep:", keep.tolist())
# keep: [0, 2]

window 1 -> classifier
window 2 -> classifier
window 3 -> classifier
...

元画像の局所領域
    |
    v receptive field
特徴マップ上の1点

候補領域を生成する
 -> proposalごとに分類とbox調整を行う

画像 -> dense detector -> boxes / classes

Images
 -> GeneralizedRCNNTransform
 -> Backbone + FPN feature maps
 -> Region Proposal Network
 -> proposals + proposal losses
 -> RoI Align
 -> box head
 -> class logits + box regression
 -> decode + score threshold + NMS
 -> detections

C5 -> P5
      upsample + C4 -> P4
                 upsample + C3 -> P3
                            ...

feature maps
 -> anchor generator
 -> RPN head
 -> objectness + bbox deltas
 -> decode
 -> clip / min-size filter
 -> pre-NMS top-k
 -> NMS
 -> post-NMS top-k proposals

sum_over_levels(H_l * W_l * anchors_per_location_l)

proposals + ground truth
 -> IoU matching
 -> positive / negative sampling
 -> RoI Align
 -> box head
 -> classification loss + box regression loss

proposals
 -> RoI Align
 -> class logits / box regression
 -> decode boxes
 -> remove background
 -> score threshold
 -> per-class NMS
 -> top detections

images: list[Tensor[C, H, W]]

model.train()
loss_dict = model(images, targets)
loss = sum(loss_dict.values())

target = {
    "boxes": boxes,    # FloatTensor[num_objects, 4], XYXY形式
    "labels": labels,  # Int64Tensor[num_objects]
}

model.eval()
with torch.inference_mode():
    predictions = model(images)

{
    "boxes": FloatTensor[num_detections, 4],
    "labels": Int64Tensor[num_detections],
    "scores": Tensor[num_detections],
}

import torch


target = {
    "boxes": torch.tensor([[10.0, 20.0, 50.0, 80.0]], dtype=torch.float32),
    "labels": torch.tensor([1], dtype=torch.int64),
}

print("boxes:", target["boxes"].shape, target["boxes"].dtype)
# boxes: torch.Size([1, 4]) torch.float32
print("labels:", target["labels"].shape, target["labels"].dtype)
# labels: torch.Size([1]) torch.int64

import torch
from torchvision.models.detection import fasterrcnn_resnet50_fpn_v2
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = fasterrcnn_resnet50_fpn_v2(
    weights=None,
    weights_backbone=None,
)

num_classes = 2  # background + pedestrian
in_features = model.roi_heads.box_predictor.cls_score.in_features  # 既存box_predictorの入力次元をそのまま流用する
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)  # 出力側だけ自分のクラス数に差し替える
model.to(device)

print("num_classes:", model.roi_heads.box_predictor.cls_score.out_features)
# num_classes: 2

import torch
from torch.utils.data import DataLoader
from torchvision.transforms import v2


def get_transform(train):
    steps = []
    if train:
        steps.append(v2.RandomHorizontalFlip(p=0.5))  # 画像とboxes/masksを同期して反転する
    steps.extend([
        v2.ToDtype(torch.float32, scale=True),
        v2.ToPureTensor(),
    ])
    return v2.Compose(steps)


def collate_fn(batch):
    return tuple(zip(*batch))  # 画像サイズや物体数が違うので、stackせずlist相当で返す


train_dataset = PennFudanDataset(
    "data/PennFudanPed",
    transforms=get_transform(train=True),
)
train_loader = DataLoader(
    train_dataset,
    batch_size=1,
    shuffle=True,
    collate_fn=collate_fn,
)

params = [parameter for parameter in model.parameters() if parameter.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.001, momentum=0.9, weight_decay=0.0005)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)

for epoch in range(num_epochs):
    model.train()
    for images, targets in train_loader:
        images = [image.to(device) for image in images]
        targets = [
            {key: value.to(device) for key, value in target.items()}
            for target in targets
        ]

        loss_dict = model(images, targets)  # 訓練モードなのでloss辞書が返る（loss_classifier、loss_box_regなど複数キー）
        losses = sum(loss for loss in loss_dict.values())  # 複数のlossを合計して1つのスカラーにする

        optimizer.zero_grad(set_to_none=True)
        losses.backward()
        optimizer.step()

    scheduler.step()

import torch
from torchvision.transforms import v2


transforms = v2.Compose([
    v2.RandomHorizontalFlip(p=0.5),
    v2.ToDtype(torch.float32, scale=True),
])

image, target

def collate_fn(batch):
    return tuple(zip(*batch))


batch = [
    ("image0", {"labels": [1]}),
    ("image1", {"labels": [2]}),
]
images, targets = collate_fn(batch)  # (image, target)のペア列を、image列とtarget列に分解する

print(images)
# ('image0', 'image1')
print(targets[0]["labels"])
# [1]

RGB image
 -> Tensor (3, H, W)
 -> model
 -> logits (C,)
 -> argmax / top-k
 -> class name

x -> conv path -> F(x) --+
|                        +-> F(x) + x
+-------- identity ------+

input (3, H, W)
 -> network
 -> logits (C, H, W)
 -> argmax over C
 -> class-id mask (H, W)
 -> paletteで可視化

high resolution -> encoder -> bottleneck -> decoder -> high resolution
                    |                         ^
                    +------ skip concat ------+

Atrous conv: カーネルのサンプル点の間隔を広げる
ASPP       : 複数のdilation rateを並列に使う
v3+        : ASPP特徴と低レベル特徴をdecoderで融合する

One-stage: image -> dense detector -> boxes / classes
Two-stage: image -> proposals -> RoI classifier / regressor

image
 -> backbone / FPN
 -> RPN anchors / objectness / deltas
 -> proposal NMS
 -> RoI Align
 -> class logits + refined boxes
 -> final NMS