簡単に自己紹介と本記事について

これが初投稿になります。地方の高専・電子情報工学科に通っている(当時)2年「さめまる」と申します。

高専の授業で"AI演習"というものがあり、自分で画像切り抜きAIを搭載したWebアプリを作ることに挑戦しました。初学者ならではの視点で書いていくので、温かい目で見ていただけると嬉しいです！

本記事では、AI部分に重きを置き、その過程と学んだこと、これから画像認識AIを作ってみたい方向けの記録です。Pythonの文法, Google Colabの使い方等はある程度理解している方向けとなりますのでご了承ください。

ディープラーニング, ニューラルネットワークの基本をおさらいしたいという方は、私がZennで書いた記事が参考になるかもしれません。
【入門・CNN】画像から背景を削除するAIを作る方法(単語, 理論をおさらい)

また、実際にこのAIを搭載したWebアプリの実装については、別の記事で紹介予定です。

当方まだ初学者なので、誤った言葉遣いや知識等あるかもしれませんが、もしあればコメントいただけると幸いです。

画像切り抜きAIの概要

早速本題です。
私はプラモデルを作るのが趣味で、「プラモデルの画像から背景を削除するAI」を作ることにしました！また、U-Netを参考に、CNNを使用した独自のモデルを設計しました。
使用技術は、画像セグメンテーション, 物体検出, 背景除去 等です。
使用環境は以下の通りです。
1. Python (Ver.3.12.9)
2. PyTorch (学習用)
3. Google Colaboratory (開発環境, T4 GPU)

開発環境としてはWindows 11 Homeを使用し、モデルの学習はGoogle Colab上で行いました。

データセットの準備

小規模な開発でしたので、過学習を狙うといいますか、1種類の、単純な背景の切り抜きに対応できることを目指しました。
今回の切り抜きの対象は、アニメ「蒼穹のファフナー」に登場するファフナー「マークニヒト」(MODEROID)です！

アノテーション

教師データは、何も加工していない「元画像」と、2値階調で量子化した「マスク画像」です。切り抜きたいファフナー本体を白, 背景を黒でIbisPaint(iPad10th)を使って塗分けしました。
このAIのゴールイメージは、
各ピクセルで、ファフナーか背景かを確率で、2値に分類できるようにすることです。

↑元画像(MODEROIDページより) ↑ 2値化画像

訓練データ用の元画像と2値化画像, 検証データ用の元画像と2値化画像を入れるディレクトリ(フォルダ)を作り、元画像と2値化画像はそれぞれ拡張子を.pngにし、名前を一緒にしました。ちなみに私は、AIのファイルと同じ階層にフォルダを作成しました。
なぜ名前を一緒にしたかというと、訓練用の画像と、その正解ラベルを紐づけしやすくするためです。

早速'.ipynb'ファイルを作成しコーディングを開始。
詳しいことは、コメントアウト参照です。

AI.ipynb

# In[1]: ライブラリ等のインポート
import torch # PyTorch
import torch.optim as optim # 最適化アルゴリズム
import torch.nn as nn # ニューラルネットワークの定義
import numpy as np # 行列演算
from google.colab import files # Google Colab上でファイルのアップロードとダウンロードの管理
import os # ファイル操作とパス管理
from PIL import Image # 画像の読み込みと操作
from torchvision import transforms # 前処理とデータ拡張
from torch.utils.data import Dataset # カスタムデータセット作成
%matplotlib inline
# 画像をセル内表示
from matplotlib import pyplot as plt # ヒートマップ作成
import datetime # 日付と時刻を表示

torch.set_printoptions(edgeitems=3) # テンソル表示は端2つのみ
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # デバイスの指定

元データが乏しく31枚しかないため、1枚につきプラス19枚の明るさ, コントラスト, 彩度等の変更を加えた計620枚を学習させました。
私は、ipynbファイルと同じ階層にフォルダを作成したので、その場合のパス指定は、/content/drive/MyDrive/Colab Notebooks/<フォルダ名>でした。

# In[2]: 前処理とデータ拡張
def Load_ConvertRGB(image_path): # 指定フォルダ内の.pngをRGBモードに変換しリストとして返す
    images = []
    for file_name in sorted(os.listdir(image_path)):
        if file_name.endswith('.png') or file_name.endswith('.PNG'):
            img = Image.open(os.path.join(image_path, file_name))
            img_rgb = img.convert('RGB')
            images.append(img_rgb)
    return images

def augment_image(image, num_augmented=19):
    transform_augment = transforms.Compose([
        transforms.ColorJitter(
            brightness=(0.8, 1.2),
            contrast=(0.8, 1.2),
            saturation=(0.8, 1.2),
            hue=(-0.1, 0.1)
        )
    ])
    augmented_images = [image]
    for _ in range(num_augmented):
        augmented_images.append(transform_augment(image))
    return augmented_images

def ToTensorWithAugmentation(images, num_augmented=19):
    augmented_tensor_list = []
    transform_to_tensor = transforms.Compose([ # 正規化
        transforms.Resize((256, 256)), # 256×256にリサイズし、
        transforms.ToTensor(), # テンソル形式に変換
    ])
    for img in images:
        augmented_images = augment_image(img, num_augmented=num_augmented)
        augmented_tensor_list.extend([transform_to_tensor(aug_img) for aug_img in augmented_images])
    return augmented_tensor_list

# データセットのパス
train_original = ''
train_threshold = ''
val_original = ''
val_threshold = ''

# 訓練データと検証データのロードと変換
train_original_images_augmented = ToTensorWithAugmentation(Load_ConvertRGB(train_original))
train_threshold_images_augmented = ToTensorWithAugmentation(Load_ConvertRGB(train_threshold))
val_original_images_augmented = ToTensorWithAugmentation(Load_ConvertRGB(val_original))
val_threshold_images_augmented = ToTensorWithAugmentation(Load_ConvertRGB(val_threshold))

アノテーションは前述のとおりディジタルペンを使ったので、2値化画像にどうしても灰色の部分が生まれてしまいました。そこで、少しでも灰色な部分はすべて白色とみなすようにしました。
また、前述のとおり「ピクセルがファフナーか、背景かの2値分類」なので、
2値化画像のサイズを(3, 256, 256)から(1, 256, 256)になるように調整しました。

# In[3]: データセットの定義
class FafnerDataset(Dataset):
  def __init__(self, original_images, threshold_images, transform=None):
    self.original_images = original_images
    self.threshold_images = threshold_images
    self.transform = transform

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

  def __getitem__(self, idx):
    original_image = self.original_images[idx]
    threshold_image = self.threshold_images[idx]
    threshold_image = (threshold_image > 0.1).float()

    if self.transform:
      original_image = self.transform(original_image)
      threshold_image = self.transform(threshold_image)

    if threshold_image.size(0) == 3:
      threshold_image = (threshold_image.sum(dim=0) > 0).float()

    return original_image, threshold_image

モデルの設計

今回形成したアーキテクチャはこちらです。

画像認識で定番のCNN(Convolutional Neural Network), 畳み込みニューラルネットワークを使用しました。
損失関数にBCE Loss, 中間層の活性化関数にReLU関数, 出力層にシグモイド関数を使用し、
畳み込み層, プーリング層, 逆(転置)畳み込み層を組み込みました。
ReLUは勾配消失問題を防ぐために、シグモイドは2値分類に適していているためです。

設計当初は、畳み込み層とプーリング層を1層とすると3層ほどのNNを組んでいたのですが、
エラーに遭遇しました。データサイズが予測値と正解値で異なる！！ と言われました。
ただ適当にNNを組んではダメだということです。

そこで、逆(転置)畳み込み層(Deconbolution)を導入し、元画像の最終的なデータサイズが2値化画像のデータサイズ(1, 256, 256)になるよう調整しました！！随時データサイズを手計算する必要があることを初めて知りました。
式は以下に記しておきます。

畳み込み層通過後の出力サイズ

(入力の1辺を$H_{in}$, 出力を$H_{out}$とする。縦横それぞれで計算が必要なので注意)
$H_{out}=\frac{H_{in}-kernel+2×padding}{stride}+1$

最大値プーリング層通過後の出力サイズ

$H_{out}=\frac{H_{in}}{kernel}$

逆畳み込み層通過後の出力サイズ

$H_{out}=(H_{in}-1)×stride-2×padding+kernel$

第○層	層の種類	画像サイズの遷移
1	畳み込み層&プーリング層	[入力]3×256×256 → 16×256×256 → 16×128×128
2	畳み込み層&プーリング層	16×128×128 → 32×128×128 → 32×64×64
3	畳み込み層	32×64×64 → 64×64×64
4	逆畳み込み層	64×64×64 → 32×128×128
5	逆畳み込み層	32×128×128 → 16×256×256
6	畳み込み層	16×256×256 → 1×256×256[出力]

具体的なフィルタ(カーネル)サイズ, パディングの値については以下のコードを参照です
nn.Moduleから継承したNNを構成。
順伝播するようforward関数で定義しました。

# In[4]: NNの定義
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(16)
        self.act1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(kernel_size=2)

        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(32)
        self.act2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(kernel_size=2)

        self.conv3 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.bn3 = nn.BatchNorm2d(64)
        self.act3 = nn.ReLU()

        self.conv4 = nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2)
        self.conv5 = nn.ConvTranspose2d(32, 16, kernel_size=2, stride=2)
        self.conv6 = nn.Conv2d(16, 1, kernel_size=3, padding=1)

        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        out1 = self.pool1(self.act1(self.bn1(self.conv1(x))))
        out2 = self.pool2(self.act2(self.bn2(self.conv2(out1))))
        out3 = self.act3(self.bn3(self.conv3(out2)))
        out = self.conv4(out3)
        out = self.conv5(out)
        out = self.sigmoid(self.conv6(out))
        return out

バッチ正規化で学習の安定化・高速化をしながら、特徴を抽出するエンコーダ部とサイズ復元のデコーダ部から構成され、最終的には確率を出力するNNを定義しました。U-Net

学習

まずは、データセットとデータローダの作成です。
訓練データはシャッフルをオンにして、バッチサイズ32の計20+1(余りも学習させる)バッチ作成しました。(1回につき2セットの学習を同時に行いました。)
うまくいかない場合はnum_workers=0にするとよいです
(num_workersは、並列処理のワーカー数を指定する値です。)

# In[5]: カスタムデータセットの作成
train_dataset = FafnerDataset(train_original_images_augmented, train_threshold_images_augmented)
val_dataset = FafnerDataset(val_original_images_augmented, val_threshold_images_augmented)

# In[6]: データローダの作成
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=32, shuffle=True, num_workers=2)
# val_loader = torch.utils.data.DataLoader(dataset=val_dataset, batch_size=4, shuffle=True, num_workers=2)

最適化アルゴリズム(オプティマイザ)には確率的勾配降下法(SGD)を、
損失関数はバイナリクロスエントロピーロス(BCE Loss)を使用し、学習率を0.01で200回学習させました。
10回ごとに損失の値を出力するように定義をし、学習させました。

# In[7]: トレーニングループの定義
def training_loop(n_epochs, optimizer, model, loss_fn, train_loader):
  for epoch in range(n_epochs + 1):
    loss_train = 0.0

    for imgs, labels in train_loader:
      imgs, labels = imgs.to(device), labels.to(device)

      labels = labels.unsqueeze(1)

      outputs = model(imgs)

      loss = loss_fn(outputs, labels)

      optimizer.zero_grad()

      loss.backward()

      optimizer.step()

      loss_train += loss.item()

    if epoch == 1 or epoch % 10 == 0:
      print("{} Epoch:{}, Training Loss:{:.4f}, {}".format(
          datetime.datetime.now(), epoch, loss_train / len(train_loader), outputs.shape
      ))

# In[8]: トレーニングループ
model = Net().to(device)
optimizer = optim.SGD(model.parameters(), lr=1e-2)
loss_fn = nn.BCELoss()

training_loop(
    n_epochs =  200,
    optimizer = optimizer,
    model = model,
    loss_fn = loss_fn,
    train_loader = train_loader
)

すると...

最終的なロスは0.0802と、大域的最適解周辺で収束したのではないかなと思います。

形式は.ptで、保存を忘れずに。

# In[9]: 保存
data_path = "./"

torch.save(model.state_dict(), data_path + "FafnerCut.pt") # ""にファイル名.pt
files.download("FafnerCut.pt")

結果

ヒートマップで精度を確認しました。(検証データを使うべきなのですが、素材が少なかったので使いませんでした。)
閾値は0.6、つまり「ファフナーであるという確率が60%以上のピクセルを赤くする」ヒートマップを作成しました。

# In[10]: ヒートマップ
def preprocess_image(image_path):
    original_img = Image.open(image_path).convert('RGB')
    original_size = original_img.size
    img_resized = original_img.resize((255, 255))
    img_tensor = torch.from_numpy(np.array(img_resized).transpose((2, 0, 1))).float() / 255.0
    img_tensor = img_tensor.unsqueeze(0)
    return original_size, img_tensor

def generate_heatmap(image_path, model, shiki):
    model.eval()
    with torch.no_grad():
        original_size, img_tensor = preprocess_image(image_path)
        img_tensor = img_tensor.to(device)
        output = model(img_tensor)

    heatmap = output.squeeze().cpu().numpy()  # 余分な次元を削除
    heatmap = (heatmap - heatmap.min()) / (heatmap.max() - heatmap.min())

    heatmap_img = plt.cm.jet(heatmap)[:, :, :3]  # カラーマップ適用
    heatmap_img = (heatmap_img * 255).astype(np.uint8)  # [0, 255] の範囲に変換
    heatmap_img = Image.fromarray(heatmap_img)  # PIL画像に変換
    heatmap_img = heatmap_img.resize(original_size)

    original_img = Image.open(image_path).convert('RGB')
    return original_img, heatmap_img

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Net()
model.load_state_dict(torch.load("ptファイルのパス", map_location=torch.device('cpu')))
model.to(device)
model.eval()

shiki = 0.6 # 閾値の設定
image_path = "画像パス"

original_img, heatmap_img = generate_heatmap(image_path, model, shiki)

plt.figure(figsize=(12, 6))

plt.subplot(1, 2, 1)
plt.title("Original Image")
plt.imshow(original_img)
plt.axis("off")

plt.subplot(1, 2, 2)
plt.title("Heatmap")
plt.imshow(heatmap_img)
plt.axis("off")

plt.show()

すると、

どうでしょうか...
背部ユニットの先端まで切り抜けていて、我ながら良い精度だと思います。

感想・問題点

私はこれまでC, C++, Processingしか触ったことがなく、PythonやWeb開発, AIの知識は0でしたが、だいたい3ヵ月ほどで単純な背景のファフナーを切り抜ける精度まで持ってこられたことに自分自身がとても驚いています。
特殊メソッドやコンストラクタの理解に苦しみました。。。
一番大変だったのはアノテーションです。1枚に10分近くかかりました...

問題点としては、まず精度です。単純な背景(単色, 境界線がくっきりしているの)なら切り抜きはそこそこといったところですが、複雑な背景になると難しいです。今回は過学習狙いでしたが、もう少し元データを増やしたり、また違う種類のプラモデルも学習させたりと、より高度なAIにしていきたいなと思いました。

AI初挑戦でどんなものを作ればいいかわからないという方は、画像認識・セグメンテーションを強くお勧めします。
ヒートマップで高精度な予測ができたときの達成感は半端ないです！！

Webアプリとして実装したものを後日投稿予定です。
ぜひそちらの記事も読んでいただけると幸いです。

最後までご覧いただき、ありがとうございました。

参考文献

「PyTorch実践入門～ディープラーニングの基礎から実装へ」, イーライ・スティーブンス(訳小川雄太郎)等, 2021年初版

編集者のGitHub

こちらにソースコードがあります。
modelディレクトリのipynbファイルがそれです。
また、READMEにその他リンクが貼ってありますのでそちらもぜひ。