Extreme Pointを使った擬似セグメンテーションラベルの生成

本記事は、2020年インターンシッププログラムで勤務された侯明程さんによる寄稿です。

はじめに

PFN2020夏期インターンに参加した侯明程です。普段は東京大学情報理工学系研究科で機械学習の研究をしています。今回のインターンではコンピュータビジョン(CV)での少数サンプルにおけるモデル学習に関するタスクをやりたいと思い、interactive object segmentationのタスクに取り組みました。

背景

一般にobject segmentationで用いられている多層のCNNの学習には大量の教師データが必要となります。その教師データの作成にはピクセルレベルでのアノテーション情報が必要であり、画像分類や物体検出といった他のCVのタスクのアノテーションと比べてもコストが高いです。またMS COCOやOpen Images のような大規模データセットでのアノテーションでは、上述した高コストのアノテーションに対して、2段階のアノテーションを行うという工夫がされています。１段階目ではfull maskのアノテーションよりも短い時間で得られる、数ピクセルのみのアノテーションとクラス情報を全部の画像に与えます。２段階目で目的のfull maskのアノテーションを行います。２段階に分けることで、１段階目が終わった時点で2段階目のアノテーションコストの見積もりができる、という利点があります。そこで、1段階目のアノテーションが終わり、２段階目のアノテーション時に点の入力から、full maskを予測するモデルがあるとアノテーションコストを大幅に削減することが期待できます。

問題設定

点の入力やbounding boxと言った、より簡単なアノテーションからピクセルレベルのfull maskのラベルを推論するモデルを構築する手法[1, 2]は既にいくつか提案されています。しかし，既存のモデルを学習データに存在しないクラスに対して適用すると性能が下がってしまいます。そこで、既に得られている、入力された点やbounding boxを基に、新しいクラスに対してfine-tuning を行うことで新しいクラスでのモデル性能を上げることができるのではないかと考えました。より具体的には、以下の2ステップから成ります。

1. 大量のfullのmaskがあるデータセットを用いて、物体の4端点を基に、対象の物体のsegmentation mask を推論するモデルを訓練する
2. その後fullのmaskが0枚または少数しかない新しいクラスに対してfine-tuniningする

この問題設定で、2の新しいクラスにおけるモデル性能を向上することが目標です。

提案手法

ナイーブに考えられるfine-tuninig の方法として、物体の4端点をラベルにし、弱教師あり学習のようにして解く手法が挙げられます[3]。しかしこの手法ですと画像全体のうち４点しか情報を渡すことができず、非常にスパースになってしまいます。さらに別の手法として4端点から生成できるbounding boxをラベルとして解く手法も考えられます。しかし下図にbounding boxと4端点でのラベルを示しましたが、ラベルとしてあまり綺麗ではありません。

そこで、4端点は物体の境界の一部を表していると考えられるので、物体の境界から物体のsegmentation maskを得る手法[4]を用いてより良いラベルが得られるのではないか、と考えました。

この手法の概要を図1に示しました。入力画像についてneural network等のモデルを用いて物体の境界の強さと向きを予測します。それを基に複数ピクセルを一つのまとまりとしたノードを構築し、複数のノードをさらにまた一つのノードにまとめて、一つの木を作成します。ノードは一つの予測されたsegmentation maskに対応し、何かしらの評価基準に基づく重みを持ち、その重みが最大または最小のノードを取得することで、最良のsegmentation maskが得られます。

ノードの重みを計算する評価基準に既に得られている4端点を使うことでより良いmaskが得られるのではないかと考えられます。そこで以下の4つの評価基準を用い、それらの線形和を取ることで最終的なノードの重みとし、この値が最小の重みのノードに対応するsegmentation maskを用いて、モデルのfine-tuningのラベルとしました。この評価基準を含めたアルゴリズムをExtreme Point Cut (EPC) と名付けました。

1. ノードのmaskから計算される4端点と実際の4端点との平均距離
2. ノードのmaskのエッジから実際の4端点との平均距離
3. ノードのmaskを囲むbounding boxが実際の4端点から計算されるbounding boxからはみ出している部分の面積の平方根
4. 実際の4端点から計算されるbounding boxがノードのmaskを囲むbounding boxよりはみ出している部分の面積の平方根

下図に各評価基準によってEPCによるキリンのmaskの良さを示しました。4つの評価基準を組み合わせることで最終的に一番良いmaskが得られるようになっています。さらにこのEPCによるマスクを用いて学習する際にはマスクがかかっている部分のみについてロスを計算するようにしました。 

実験

提案手法の有効性を示すために以下の実験を行いました。まず用いたデータセットは、MS COCOという一般的な画像データセットとDRIONSという医療画像データセットです。MS COCOには80クラスのsegmentation アノテーションがついており、これらのうち5クラスをモデルの訓練に使い、残りの75クラスのうち評価指標であるIoUが悪かった3つのクラスに対してfine-tuning を行いました。モデル訓練時に使われていないクラスのMS COCOのデータセットをMS COCO (novel) と呼ぶことにします。また、MS COCOの80クラスのうち60クラスを用いてモデルを訓練し、DRIONSでもfine-tuning を行いました。まとめるとfine-tuningのシナリオは以下のようになります。

1. MS COCO -> MS COCO (novel)
   1. MS COCO 5 class でモデルを訓練
   2. 他の75クラスに対してIOUが低かった下位3つのクラスに対してfine-tuning
2. MS COCO -> DRIONS
   1. MSCOCO 60 class でモデルを訓練
   2. 医療画像DRIONS(http://www.ia.uned.es/~ejcarmona/DRIONS-DB.html)でfine-tuning

MS COCOは一般的な画像データセットで、80 クラスのsegmentation annotation があります。 DRIONSは医療画像データセットで、網膜に関するsegmentation annotation が提供されています。モデルはVGG19をencoderにしたUNetにskip-connectionのついたものを用いました。また、fine-tuningの際には新クラスにoverfitするのを防ぐために訓練データ(aのデータ)と新クラスのデータ(bのデータ)を5:1で混ぜてモデルの訓練を行いました。EPCによるfine-tuningに加えて、真のラベルを用いたfine-tuningも行いました。これはEPCのラベルが不完全であるので、完全なラベルを用いてfine-tuningによる精度向上の上限を見積もるためです。また、EPCのmask自体によるスコアも示しました。

MS COCO -> MS COCO (novel)の結果

MS COCOの5クラスでモデルを訓練し、残りの75クラスのうちスコアが下位3クラスのものについての評価結果を表に示しました。評価指標はobject segmentationで用いられるIoUを用いました。結果より、EPCによるfine-tuningではfine-tuningなしよりは精度が上がりますが、さほど変わらない、という結果になりました。真のラベルを用いたfine-tuningでも、2~4%程度の上昇と、大きく上昇することはありませんでした。上昇幅が小さい原因としては、モデルが見たことないクラスだから識別できないのではなく、形状が細いものが多く、segmentation自体が難しいのでスコアが低いのではないか、と考えられます。

MS COCO -> DRIONS の結果

DRIONSでは視神経頭(optic nerve head)というクラスのsegmentationを行います。下の表にFrozen model、EPCによるfine-tuning、真のラベルによるfine-tuning、EPCmaskスコアを示しました

EPCのmask自体のIoUが30.4と低いにも関わらず、EPCによるfine-tuningでは16%程度も精度が向上し、真のラベルを用いたfine-tuningと近い精度になりました。

上図はモデル予測の可視化結果で、1列目は入力画像、2列目はFrozenモデルの予測結果、3列目はEPCでfine-tuningしたモデルの予測結果、4列目は真のラベルであり、EPCでfine-tuningしたモデルではかなり正解ラベルに近い結果となっています。さらに提案手法の章で述べた、4端点を用いた学習とbounding boxを用いた学習の性能も下表に示しました。また右図は各手法で訓練したモデルの予測結果を可視化したものです。

表及び可視化結果より、EPCでのfine-tuningが最も性能が良く、またboudning boxと4端点では予測maskが真のラベルからかけ離れたものになっています。次に新しいクラスの真のラベルの枚数を0枚から変えていったときにスコアがどう変わるかを調べました。

表から、真のラベル付きデータの枚数を増やせば性能も上がっていきますが、既に見たように真のラベルデータでfine-tuningした時とEPCでfine-tuningした時のスコアの差が小さかったので上がり幅はそれほど大きくはありません。また上図左は200iterationでfine-tuningした際の各iteration毎のスコアで、上図右はそれを一部拡大したもので、真のラベルの枚数が変わっても性能や学習曲線はそれほど変わっていないことがわかります。

提案手法の限界と展望

上図は DRIONS データセットでのEPC Maskを可視化したもので、1列目が画像、2~6列目がEPC Weightが小さい順にtop5を左から並べたもので、一番右の7列目が真のラベルです。この可視化結果から、EPC Maskは網膜の線のエッジに影響されたMaskを出してしまっていることがわかります。またEPC Maskは総じてPrecisionが高いことが言えます。このことから、今はEPC Weightが最も小さいノードを取ってきていますが、複数ノードの組み合わせでスコアを算出して最適なノードを選択すればより良いマスクが得られると考えられます。

まとめ＆感想

Neural Networkに大量のラベルが必要ですが、そのラベル生成にDeep Learning 以外の手法を使い、擬似ラベルを生成することが結構有効なことがわかりました。今年のインターンはリモートでの開催となり、期間も例年より短い約一ヶ月でしたが、様々な方によるサポートがあり、有意義なインターンにすることができました。特にメンターの二井谷さん、高木さんには毎日ミーティングをして下さり、オンラインであるにも関わらずコミュニケーションが密に取れて非常に助かりました。僕自身研究の経験がなく、取り組んだテーマや目標達成の難しさを実感しましたが、メンターさん達の支え合って、なんとか目に見える形になりました。また、インターンを経て、仮説を立てそれを検証するというサイクルの速さがキーになると実感し、基本ではありますがサーベイの大事さも学びました。改めて、まだまだ学ぶことが多いと実感したインターンでした。PFNで学んだことを忘れずにこれからも励んでいきたいと思います。最後になりますが、メンターの二井谷さん、高木さんには毎日ミーティングをし、上手くいかない時もアドバイスをしてくだり、多大なサポートをして頂き、本当にありがとうございました。

参考文献

[1] K.K. Maninis, S. Caelles, J. Pont-Tuset, and L. Van Gool. Deep Extreme Cut: From Extreme Points to Object Segmentation, Computer Vision and Pattern Recognition (CVPR), 2018.

[2] Z. Shiyin, L. Jun Hao , W. Yunchao , W. Shikui, Z. Yao, Interactive Object Segmentation with Inside-Outside Guidance , Computer Vision and Pattern Recognition (CVPR), 2020.

[3] A. Bearman, O. Russakovsky, V. Ferrari, and L. Fei-Fei. What’s the point: Semantic segmentation with point supervision. arXiv:1506.02106v4 [cs.CV], 2015.

[4] P. Trung ,D. Thanh-Toan, S. Niko, R. Ian, SceneCut: Joint Geometric and Object Segmentation for Indoor Scenes,  CoRR, abs/1709.07158, 2017.