エムスリーテックブログ

エムスリー(m3)のエンジニア・開発メンバーによる技術ブログです

CycleGANで効果を発揮!画像の構造を抽出する特徴量"MIND"を実装してみる

 初めまして!2019年8月中旬からエムスリー エンジニアリングG AIチームで10日間インターンに参加した三澤です。インターンでは「CycleGANを用いてモダリティ(CT, MRI, PETなどの画像撮影装置)の違う画像の変換に関する手法」に関する論文について、Surveyと実装をしました。CycleGANというのは「夏の風景画と冬の風景画」や「写真と絵画」などといった画像のスタイルを変換するDeep Laerningの手法です。医療画像診断においては、「MRI画像とCT画像」といった、モダリティの異なる画像の変換にCycleGANが使えるのではないかと研究されています。

 本記事では、CycleGANに応用ができるMINDと呼ばれる特徴量(正確には記述子)が非常に興味深かったので、紹介して実装していこうと思います。

MINDの登場とCycleGAN

 2012年に、医療用画像の複数のモダリティ間の画像のレジストレーション(位置合わせ)という文脈でMIND(Modality Independent Neighborhood Descriptor) [1]が提案されました。これは組織などに依らず画像の構造を抽出しようというものです。
 このMINDがCycleGANにも使えるという話が最近報告されています。医療画像におけるCycleGANでは、例えばMRI画像とCT画像を変換しようとしたとき、教師データが「同じ人の同じ部位のMRI画像とCT画像である」(=ペアである)と良いのですが、 そのようなデータを大量に用意するのは工数もかかり大変な作業です。また教師データがペアでは無い場合、同じ輪郭を持った画像が生成されるとは限らないという点が課題でした。 今回注目する論文[2]は、ペアではない医療画像のモダリティの変換において、「CycleGANにMINDという特徴量に基づく損失(Loss)を加える」と教師データがペアである場合にヒケを取らない結果が出た、と主張するものです。
 ここで、もし構造を抽出するMINDが機能すれば、教師データをペアになるようにデータを集めたりレジストレーションしたりする工程が省かれるということが考えられます。加えてCycleGANでモダリティの変換ができれば画像診断におけるデータ拡張でも今後使えますし、またMINDそれ自身を特徴量として用いることができます。
  今回はMINDの計算方法が結構複雑なのでそれを解説した後、Keras(Python)で実装をしていこうと思います。

MINDの計算方法

 I を画像を表す二次元配列とします。 x はピクセルの位置を表す順序対のようなものとします。例えば  x=(1, 2) の場合、 I(x) I (1, 2) の位置の値を表すとします。 またMINDを計算するにあたって I の任意の点  x において、  x 周りの正方形領域(以下ではnon-local regionと呼びます)を考えます。また、R_{nl}と書いてnon-local regionを表すとします。以下の画像はnon-local regionが  5\times 5の場合です。 x を中心に配置したいのでnon-local regionの一辺のピクセル数は奇数を用います。これはMINDのパラメータです。

f:id:somisawa:20190829150749p:plain
x のnon-local region

 I 上の任意のピクセル  x において、その点におけるMINDはnon-local regionと同じshapeの二次元配列になっています。点  x における  \alpha \in R_{nl} の位置での MINDの値は次のように定義されます:


\begin{aligned}
F_x^{(\alpha)}(I) := \frac{1}{Z} \exp{\left(-\frac{D_{\mathcal{P}}(I, x, x+\alpha)}{V(I, x)}\right)}
\end{aligned}

 ここで、 ZF _ xの要素の最大値が1になるように規格化する因子です(この後の定義を見ればわかりますが  D_\mathcal{P} の各成分は非負で  D _ \mathcal{P}| _ {\alpha = 0} =0 なので常に Z=1 だと思いますが論文にあったので念の為 )。あとで説明するように、MINDは平行移動させたときの相関のようなものを表しています。 D _ {\mathcal{P}}(I, x, x+\alpha)  V(I, x) は以下に定義を書きます。

 D _ {\mathcal{P}}(I, x, x+\alpha)は次のように定義されます:


\begin{aligned}
D _ {\mathcal{P}}(I, x, x+\alpha) := \sum _ {p\in \mathcal{P}} \left(I(x+p) - I(x+\alpha +p)\right)^2
\end{aligned}

 ここで \mathcal{P} はpatchと呼び、上の式では  x とnon-local region内の点  x+\alpha との値の二乗誤差をpatch上で足しています。ここでpatchのサイズもnon-local regionと同様MINDのパラメータです。

f:id:somisawa:20190829154243p:plain
patchが 3\times3 の場合のD _ \mathcal{P}の図

 勘のいい方はお気づきかと思いますが、二次元の添字をぐりぐり動かして足し合わせているのでこれは二次元畳み込みで表現されます。二次元畳み込みについて詳しくは実装パートで触れます。

 そして  V(I, x) D _ {\mathcal{P}}(I, x, x+\alpha) を用いて次のように定義されます:


\begin{aligned}
V(I, x) := \frac{1}{4}\sum _ {n\in \mathcal{N}} D _ {\mathcal{P}}(I, x, x+n)
\end{aligned}

 ここで \mathcal{N} x に隣接する上下左右の4ピクセルの集合を表します。

 以上でMINDが定義されました。 D _ {\mathcal{P}}(I, x, x+\alpha) はそれなりに複雑に思う人もいるかもしれないので画像を使って定性的に説明しようと思います。以下に何故構造が抽出できるのかのイメージを示します:

f:id:somisawa:20190829161537p:plain
xのnon-local region
f:id:somisawa:20190829161613p:plain
ピクセル x におけるMIND  F _ x

実装

実行環境

OS: macOS Mojave 10.14.6
CPU: 2.3 GHz Intel Core i5
メモリ: 16 GB 2133 MHz LPDDR3
Python: 3.7.4

requirements(抜粋)

tensorflow==1.14.0
Keras==2.2.5

 まず必要なライブラリをインポートします。

import numpy as np # バージョンによってはTensorFlowのwarningが出る

#KerasはバックエンドとしてTensorFlowを指定してください。
from tensorflow import keras
import keras.backend as K
from keras.layers import add, subtract, multiply, Reshape, Lambda

 D_\mathcal{P} の計算は二次元畳み込みになっていると述べましたが、具体的には全ての要素が1のpatchと同じshapeの二次元配列を  (I - I' _ \alpha) ^ 2 という画像と畳み込めば良いです(返す二次元配列はshapeが変わらないようにpaddingします)。ここで  I' _ \alpha I \alpha だけ平行移動させた画像です。以下に平行移動をさせる関数と D _ \mathcal{P} を計算する関数を示します:

# デフォルトはグレースケールを想定
def make_kernel(x, y, nl_region_size, chs=1):
    # 平行移動をさせるカーネルを返す
    kernel = np.zeros((nl_region_size, nl_region_size, chs, 1))
    kernel[x, y, :, :] = 1
    return kernel

def shift_img(img, x, y, nl_region_size, chs=1):
    # 画像を平行移動させる。Affine Transformだと遅いので二次元畳み込みで計算する。
    kernel = make_kernel(x, y, nl_region_size, chs)
    kernel_k = K.constant(kernel)
    output = Lambda(lambda img: K.conv2d(img, kernel=kernel_k,
                                         padding="same", strides=1, dilation_rate=1))(img)
    return output

def calc_dp(img, simg, patch_size, chs=1): # 動的計画法ではない
    C = K.ones((patch_size, patch_size, chs, 1))  # TODO: Gaussian Kernel -> 中心の値が大きくてpatchの端に行くほど小さいカーネル
    diff = subtract([img, simg])
    norm = multiply([diff, diff])
    output = Lambda(lambda x: K.conv2d(x, kernel=C, padding="same"))(norm)
    return output

 ここで、平行移動は二次元のデルタ関数ライクな配列との畳み込みによって表されることを使うと(二次元畳み込みはDeep Learningではよく使うためそれなりにアルゴリズムが最適化されていそうなので)、愚直にAffine Transformをやるより早く処理ができそうです。入力される画像は(batch_size, height, width, channel)という順序であることを想定しています。

 これで D_\mathcal{P} が計算できるようになったのでMINDを計算する関数を実装します。

# グレースケールを想定
# 入力する画像はtensorで、(n, height, width, channel)。nはNoneでも可
def MIND(img, patch_size, nl_region_size, batch_size):
    # patch_sizeとnl_region_sizeは奇数
    assert patch_size % 2 == 1
    assert nl_region_size % 2 == 1
    
    n_imgs = batch_size
    range_alpha = (nl_region_size - 1)//2

    # 平行移動させる二次元配列
    alphas = [[i, j] for i in range(nl_region_size) for j in range(nl_region_size)]

    Fx, V = None, None
    flag_first_V = True

    for i, alpha in enumerate(alphas):
        # 平行移動させた画像を作る
        simg = shift_img(img, alpha[0], alpha[1], nl_region_size)

        # Dpを計算
        dp = calc_dp(img, simg, patch_size)
        dp = Reshape((img.shape[1], img.shape[2], 1, 1))(dp)

        # Fxを積み重ねていく
        Fx = K.concatenate([Fx, dp], axis=4) if i else dp

        # Vを計算する
        if (abs(alpha[0]-range_alpha)*abs(alpha[1]-range_alpha) == 0) and (abs(alpha[0]-range_alpha)+abs(alpha[1]-range_alpha) == 1):  # 隣接かどうか
            if flag_first_V:
                V = dp
                flag_first_V = False
            else:
                V = add([V, dp])

    V = Lambda(lambda x: x/4)(V)
    rev_V = Lambda(lambda x: 1/(x+1e-2))(V) # <- 1/V
    
    Fx = Lambda(lambda x: -x)(Fx)
    Fx = K.exp(multiply([Fx, rev_V]))
    
    rev_Fx_ma = Lambda(lambda x: 1/x)(K.expand_dims(K.max(Fx, axis=4), axis=-1)) # 規格化のために最大値の逆数を計算
    Fx = multiply([Fx, rev_Fx_ma]) # 規格化
    Fx = K.reshape(Fx, (n_imgs, img.shape[1], img.shape[2], nl_region_size, nl_region_size)) # 整形

    # Fxは(n, height, width, nl_region_size, nl_region_size)になっている。axis=1, 2(0-indexed)にxの位置を指定するとその点でのMINDを表す。
    return Fx

結果

 今回用いた画像データはこちら(IXI Dataset)[3]で配布されている頭部MRI画像を用いました。こんな画像です:

f:id:somisawa:20190829183527p:plain
頭部MRI画像

 この画像でnon-local regionを  9\times 9 としていい感じの  x を見つけてそのnon-local regionを図示すると次のようになります:

f:id:somisawa:20190829204300p:plain
 x でのnon-local region

 この  x でのMINDを求めてみましょう。patch shapeは  7 \times 7としました。

f:id:somisawa:20190829190604p:plain
MIND

 うまく構造が抽出できているように見えます!

 最後に、画像のモダリティが違った状況で、構造が同じなのにMINDが全く違う場合本末転倒なのでそこは検証する必要があります。その検証にあたって、本来ならMRIとCTなどのペアの画像を使って検証できれば良いのですが、残念なことにそのような画像でパブリックなものが見あたりませんでした。そこで今回は、先ほどの画像を適当な非線形フィルタに通したあと強い白色ノイズを加算した画像を使って比較することにします。

f:id:somisawa:20190829185624p:plain
MRI画像(左)と加工した画像(右)

 先ほどのnon-local regionに着目してみると以下のようになります:

f:id:somisawa:20190829204734p:plain
MRI画像(上)と加工した画像(下)の  x のnon-local region

 この2枚に対してMINDを計算させると次のようになります:

f:id:somisawa:20190829190527p:plain
MRI画像(左)と加工した画像(右)の  x のMIND

 図のように双方で同じような画像ができているのが見て取れます。今回は強めにノイズをかけて組織がわかりづらいように調整したのですがそれでも構造が見えているので良いと思います。

CycleGANで使うには

 CycleGANで使う場合、2つのドメインでのGeneratorの生成画像と元の画像のMINDの平均二乗誤差(MSE)などをLossに加えると良いでしょう。

最後に

 今回は画像の構造を抽出するMINDを実装しました。自分はまだ学部3年で自分から論文を読んで実装する機会がなかったため、インターンではそのような機会を与えていただき、さらにアウトプットまでさせていただけることになり非常に良い経験となりました。

 エムスリーでは随時インターンを募集しているのでこの記事を読んで興味を持っていただいた方は下記リンクを是非チェックしてください。 jobs.m3.com


  1. Heinrich, M.P., et al.: MIND: Modality independent neighbourhood descriptor for multi-modal deformable registration. Med. Image Anal. 16(7), 1423–1435 (2012)

  2. Yang, H. et al. Unpaired brain mr-to-ct synthesis using a structure-constrained cyclegan. In Deep Learning in Medical Image Analysis and Multimodal Learning for Clinical Decision Support, 174–182 (Springer, 2018).

  3. IXI Dataset, made available under the Creative Commons CC BY-SA 3.0 license