第3回：PF-AFNの実装 — 試着エンジンのコードを読む

連載
AI Machine Learning Python GAN Computer Vision

はじめに — コードで理解する仮想試着

前回はGANの原理と、汎用的な画像生成技術がどのように仮想試着タスクに適用されるかを解説しました。今回はいよいよ、META FITの中核を担う試着エンジン --- PF-AF（Parser-Free Appearance Flow Network）の実装に踏み込みます。

PF-AFNの「Parser-Free」という名称には重要な意味があります。従来のVTON手法の多くは、人物画像のセグメンテーション（Human Parsing）を前処理として必要としていました。頭、胴体、腕、脚といった体の部位を事前に分割し、それをモデルへの入力として与える必要があったのです。PF-AFNはこの前処理を不要にし、人物画像と衣服画像の2枚だけで試着結果を生成できるアーキテクチャです。

この記事では、実際のソースコードを追いながら、推論パイプラインの全体像から各モジュールの内部構造まで、技術的な詳細を一つひとつ紐解いていきます。

2段階パイプラインの全体像

PF-AFNの推論パイプラインは、明確に分離された2つのステージで構成されています。

  1. Stage 1 --- AFWM（Appearance Flow Warping Module）: 学習されたオプティカルフローを用いて、平置きの衣服画像を人体の形状に合わせてワーピング（幾何学的変形）します。
  2. Stage 2 --- ResUnetジェネレータ: ワーピングされた衣服を人物画像に自然に合成（コンポジット）し、最終的な試着画像を生成します。

test.pyから抜粋した推論コードで、この2段階の流れを確認してみましょう。

# Stage 1: 衣服のワーピング
warped_cloth, last_flow = warp_model(real_image, clothes)

# Stage 2: コンポジット生成
gen_output = gen_model(torch.cat([real_image, warped_cloth, warped_edge], 1))
p_rendered, m_composite = torch.split(gen_output, [3, 1], 1)
p_rendered = torch.tanh(p_rendered)
m_composite = torch.sigmoid(m_composite)

# 最終合成
p_tryon = warped_cloth * m_composite + p_rendered * (1 - m_composite)

Stage 1が衣服の「形を合わせる」処理、Stage 2が「自然に合成する」処理です。最終行の合成式は、コンポジットマスク m_composite を用いたブレンディングです。マスクの値が1に近い領域ではワーピング済み衣服がそのまま使われ、0に近い領域ではジェネレータが描いたレンダリング結果が使われます。この仕組みにより、衣服のテクスチャをできるだけ保ちつつ、肌や背景との境界を自然に仕上げることができます。

Stage 1 --- AFWMアーキテクチャ

Stage 1のAFWM（Appearance Flow Warping Module）は、3つのサブモジュールから構成されています。afwm.pyのコードを追いながら、それぞれの役割を解説します。

FeatureEncoder --- マルチスケール特徴抽出

FeatureEncoderは、人物画像と衣服画像の両方に使われる特徴抽出器です。5つの畳み込み層で構成され、各層がチャンネル数を増やしながら空間解像度を半分にしていきます。

# チャンネル数の増加パターン
channels = [64, 128, 256, 256, 256]
# 各層の構成: Conv2d → BatchNorm → ReLU, stride=2でダウンサンプリング

各層の出力が保存され、5つの異なる解像度の特徴マップが得られます。最も粗い特徴マップは入力の1/32の解像度、最も細かいものは1/2の解像度です。このマルチスケールの特徴表現が、後段のフロー推定で「大まかな位置合わせ」から「細部の調整」までを段階的に行うための土台になります。

RefinePyramid --- Feature Pyramid Network

RefinePyramidは、FeatureEncoderが出力した5つのスケールの特徴マップを統合し、整ったマルチスケール特徴を作り上げるモジュールです。物体検出の分野で広く使われるFeature Pyramid Network（FPN）の考え方を取り入れています。

# 5つのエンコーダ特徴からラテラル接続を構築
# 粗い特徴をアップサンプリングして細かい特徴に加算
for i in range(4):
    upsampled = F.interpolate(coarse_feature, scale_factor=2)
    refined = lateral_conv(fine_feature) + upsampled

ラテラル接続の役割は、高解像度の特徴が持つ空間的な細かさと、低解像度の特徴が持つ意味的な情報を、各スケールの特徴に反映させることです。これにより、粗いレベルでは「衣服全体がどの位置に来るべきか」、細かいレベルでは「襟元や裾のラインをどう合わせるか」を、それぞれ適切に判断できる特徴表現が得られます。

AFlowNet --- コアイノベーション

AFlowNet（Appearance Flow Network）は、PF-AFNのもっとも重要なモジュールです。RefinePyramidが出力したマルチスケール特徴を使って、**粗から細へ（Coarse-to-Fine）**のフロー推定を行います。

5つのピラミッドレベルを、最も粗いレベルから最も細かいレベルへと順番に処理していきます。各レベルでの処理は以下の流れです。

  1. 人物特徴と衣服特徴の間の相関マップを計算する（カスタムCUDAカーネルを使用）
  2. 相関マップを人物特徴と現在のフロー推定値に結合する
  3. 畳み込み層でフロー残差（現在のレベルで追加すべき修正量）を予測する
  4. 前の（粗い）レベルからアップサンプリングしたフローに残差を加算する
  5. 更新されたフローを使って衣服特徴をワーピングし、次のレベルの入力とする

この粗から細へのアプローチの利点は明快です。最も粗いレベルでは、衣服全体の大まかな位置合わせを行います。この段階ではピクセル単位の精度は不要で、「肩の位置」「胴体の幅」といった大まかな対応関係をつかめれば十分です。レベルが細かくなるにつれて、襟の形状、袖口のライン、裾のカーブといった細部の位置合わせが段階的に精密になっていきます。

最終的な出力は、密なフローフィールド --- つまり、衣服画像の各ピクセルが人物画像上のどの位置に配置されるべきかを示す2次元ベクトルの集合です。このフローフィールドに基づいてグリッドサンプリングを行うことで、衣服画像が人物の体型に合わせて滑らかに変形されます。

CUDA相関カーネル --- パフォーマンスの要

AFlowNetの各レベルで計算される「相関マップ」は、このモデルの性能を左右する重要な処理です。しかし同時に、計算量が非常に大きいため、パフォーマンス上のボトルネックにもなり得ます。PF-AFNでは、この処理をCuPyを介したカスタムCUDAカーネルで実装しています。

correlation.pyには、2つのCUDAカーネルが定義されています。

# 相関計算のカスタムCUDAカーネル
kernel_Correlation_rearrange = '''
    extern "C" __global__ void kernel(...) {
        // 効率的な相関計算のための特徴マップ再配置
        // (batch, channel, height, width)から
        // パディング済みレイアウトへ変換
    }
'''

kernel_Correlation_updateOutput = '''
    extern "C" __global__ void kernel(...) {
        // 2つの特徴マップ間の相関を計算
        // 各位置で、カーネルサイズで定義される
        // ローカル近傍内のチャンネル方向内積を計算
    }
'''

相関計算が行っていることを直感的に説明すると、「人物画像の各位置に対して、衣服画像のどの位置が最も特徴的に類似しているか」を探索する処理です。各空間位置について、定義されたカーネルサイズの局所近傍内でチャンネル方向の内積を計算し、類似度マップを生成します。

計算量は O(H×W×K×K×C)O(H \times W \times K \times K \times C) --- 画像の高さ、幅、カーネルサイズ、チャンネル数の積に比例します。256x192ピクセルの入力画像に対して5つのスケールレベルで繰り返し実行されるため、純粋なPython実装では実用的な速度が出ません。CUDAによるGPU並列計算が不可欠であり、これがインフラ面での大きな要件になっています。

具体的には、CuPyライブラリとNVIDIA GPUが必須です。CuPyはPythonからCUDAカーネルを直接書いて実行できるライブラリで、NumPy互換のAPIを提供しつつGPU上での演算を可能にします。このGPU依存が、開発環境にDockerコンテナ（NVIDIA NGC）を採用した主な理由の一つでもあります。

Stage 2 --- ResUnetジェネレータ

Stage 1でワーピングされた衣服を、そのまま人物画像に貼り付けるだけでは自然な結果になりません。衣服と肌の境界、影の落ち方、衣服を透かして見える体のラインなど、見た目の整合性を保つ処理が必要です。この役割を担うのが、Stage 2のResUnetジェネレータです。

アーキテクチャの構成

networks.pyで定義されているResUnetジェネレータの基本構成は以下の通りです。

ResUnetGenerator(input_nc=7, output_nc=4, num_downs=5, ngf=64)
  • 入力: 7チャンネル = 人物画像（3ch） + ワーピング済み衣服（3ch） + ワーピング済みエッジマップ（1ch）
  • 出力: 4チャンネル = レンダリング人物画像（3ch） + コンポジットマスク（1ch）
  • ダウンサンプリング: 5段、チャンネル数は 64 → 128 → 256 → 512 → 512

名前の通り、U-Netの構造にスキップ接続を持ちつつ、各レベルに**残差ブロック（Residual Block）**を組み込んだ設計です。U-Netのスキップ接続はエンコーダとデコーダの対応するレベルを直結し、空間的な詳細情報の損失を防ぎます。残差ブロックは各レベルでの特徴表現をより豊かにします。

ResidualBlock --- 学習の安定化

残差ブロックの実装は、ResNetで提案された残差接続の考え方をそのまま踏襲しています。

class ResidualBlock(nn.Module):
    def __init__(self, in_features):
        self.block = nn.Sequential(
            nn.ReflectionPad2d(1),
            nn.Conv2d(in_features, in_features, 3),
            nn.InstanceNorm2d(in_features),
            nn.ReLU(inplace=True),
            nn.ReflectionPad2d(1),
            nn.Conv2d(in_features, in_features, 3),
            nn.InstanceNorm2d(in_features)
        )

    def forward(self, x):
        return x + self.block(x)  # 残差接続

return x + self.block(x) --- この一行が残差接続の本質です。ブロック内の畳み込み層は入力そのものではなく、入力に対する「修正量」（残差）を学習します。これにより、深いネットワークでも勾配が安定して伝播し、学習が効率的に進みます。

InstanceNorm2dを使っている点にも注目です。BatchNormではなくInstanceNormにすることで、各サンプルごとに正規化が行われます。スタイル変換や画像生成タスクでは、バッチ内の統計量に依存しないInstanceNormの方が良い結果を出す傾向があり、ここでもその知見が活かされています。ReflectionPad2dは、境界でのアーティファクトを軽減するためのパディング手法です。

コンポジットマスクの役割

ResUnetジェネレータの出力4チャンネルのうち、最後の1チャンネルがコンポジットマスクです。このマスクが最終的な画像合成の品質を大きく左右します。

p_rendered = torch.tanh(p_rendered)      # [-1, 1]の範囲
m_composite = torch.sigmoid(m_composite)  # [0, 1]の範囲

# 最終合成
p_tryon = warped_cloth * m_composite + p_rendered * (1 - m_composite)

sigmoid関数を通したマスク値は0から1の範囲を取り、各ピクセルにおけるワーピング衣服とレンダリング結果のブレンディング比率を決定します。

  • マスク値が1に近い領域（主に衣服の中央部分）: ワーピングされた衣服をそのまま使用します。衣服のテクスチャ、パターン、色味を忠実に保持するために重要です。
  • マスク値が0に近い領域（肌、背景、元の衣服を除去した部分）: ジェネレータがレンダリングした画像を使用します。肌の色味や背景の一貫性を保ちます。
  • 中間値の領域（衣服と肌の境界）: 両者がスムーズにブレンドされ、自然な遷移を実現します。

データパイプラインとインフラ

入力仕様

aligned_dataset_test.pyで定義されている入力データの仕様をまとめます。

項目仕様
画像サイズ256 x 192ピクセル
人物画像全身写真（正面）
衣服画像フラットレイ（平置き）商品写真
エッジマップ衣服のアウトライン（前処理済み）
前処理幅192にスケール → [256, 192]にリサイズ → [-1, 1]に正規化

256x192という解像度は、計算コストと生成品質のバランスを考慮した選択です。GPUメモリの制約の中で実用的な速度を維持しつつ、衣服のテクスチャを判別できる最低限の解像度を確保しています。

Docker環境

PF-AFNの実行環境は、NVIDIA NGC（NVIDIA GPU Cloud）のDockerコンテナをベースに構築されています。

# ベースイメージ: NVIDIA公式PyTorchコンテナ
FROM nvcr.io/nvidia/pytorch:18.04-py3
# PyTorch 1.2.0 + CUDA 9.2
# CuPy 6.0.0（カスタムCUDAカーネル用）

NGCコンテナを採用した理由は、CUDAドライバ、cuDNN、PyTorch、CuPyといった複雑な依存関係をまとめて解決できるからです。特にCuPyのバージョンとCUDAバージョンの整合性を手動で合わせるのは大変で、動作確認済みのコンテナイメージを使うことで環境構築の手間とリスクを大きく減らせます。

ただし、このインフラ要件 --- Docker、NVIDIA GPU、大容量GPUメモリ --- は、本番環境への展開を考えるうえで大きな制約です。消費者向けのスマートフォンアプリとして提供するには、クラウド上のGPUサーバーで推論を実行し、結果をアプリに返す構成が必要になります。

PF-AFNからPASTA-GAN++への進化

META FITでは、PF-AFNに加えてPASTA-GAN++も試着エンジンとして評価・採用しました。両者の違いを整理しておきます。

PF-AFNは「Parser-Free」の名の通り、Human Parsingを入力に必要とせず、人物画像と衣服画像の2枚だけで動作します。このシンプルさは強みですが、体の部位に関する情報がないぶん、複雑なポーズや体型への対応には限界がありました。

PASTA-GAN++は、OpenPoseによる姿勢推定とGraphonomyによるHuman Parsingを補助入力として追加することで、より精密な衣服配置を実現しています。

項目PF-AFNPASTA-GAN++
入力人物画像 + 衣服画像人物画像 + 衣服画像 + ポーズ + パーシング
補助モデルなしOpenPose、Graphonomy
学習データVITONDeepFashion
対応モード上半身full（上下）、upper、lower

PASTA-GAN++の設定では、network-snapshot-004408.pkl（GANの学習済み重み）、body_pose_model.pth（OpenPose）、inference.pth（Graphonomy）の3つのモデルファイルが必要です。

PF-AFNで確立された「オプティカルフローによる衣服ワーピング + ジェネレータによるコンポジット」という基本パイプラインは、PASTA-GAN++でも踏襲されています。PF-AFNのコードをしっかり理解しておくことが、PASTA-GAN++の動作原理を把握するうえでも大切な基盤になります。

次回予告

今回は、PF-AFNの実装をコードレベルで解説しました。FeatureEncoderによるマルチスケール特徴抽出、RefinePyramidによる特徴統合、AFlowNetの粗から細へのフロー推定、CUDA相関カーネルによる高速化、そしてResUnetジェネレータによるコンポジット生成。これらが連携して、2枚の画像から試着結果を生成するパイプラインを構成しています。

次回は、試着エンジンの周辺を支える技術群に焦点を当てます。OpenPoseによる骨格検出、Graphonomyによる人体パーシング、そして独自に開発した自動採寸アルゴリズム。さらに、PiFuを用いた2D写真から3Dメッシュへの再構成の試みと、TensorFlow.jsによるブラウザベースのプロトタイプについても解説します。

META FIT シリーズ：

この記事をシェア

X Facebook LinkedIn

関連記事

第1回：毎日の献立問題をデータで解く
AIMachine LearningPython

第1回：毎日の献立問題をデータで解く
第2回：栄養ベクトルとコサイン類似度で「同じ栄養、違う料理」を実現する
AIMachine LearningPython

第2回：栄養ベクトルとコサイン類似度で「同じ栄養、違う料理」を実現する
第3回：LSTMで「飽きの来ない献立」を時系列予測する
AIMachine LearningPython

第3回：LSTMで「飽きの来ない献立」を時系列予測する