微分可能画像処理ライブラリ Kornia への招待

はじめに

本エントリは Sansan Advent Calendar 2019 第9日目の記事です． R&D系の話題に限らず，多様な話題が投稿されますのでぜひウォッチしていただけると幸いです 🙏

adventar.org

TL; DR

Korniaはイイぞ

PyTorch + Computer Vision

PyTorchといえば，言わずと知れた機械学習フレームワークですよね． 検索数はここ1-2年で一気にChainerを抜き去り*1，TensorFlowにも肉薄しています． TensorFlowはプロダクションにおける人気が絶大ということもあって，研究サイドにフォーカスするとPyTorch人気はより一層強調されます*2．

Computer Vision界隈においてPyTorchが人気な理由の一つとして，torchvision の存在があります． torchvisionは画像データローダにおける前処理や学習済みモデルを提供するライブラリです． 特に前処理が簡単に書けるのが良いところで，込み入ったことをしない限り，次のようにスッキリ書けます．

from pathlib import Path

from PIL import Image
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms


class OleOleDataset(Dataset):
    def __init__(self, data_dir):
        self.filenames = list(Path(data_dir).glob('*.png'))
        self.transforms = transforms.Compose([  # 前処置を積んでく
            transforms.RandomResizedCrop(224),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
        ])

    def __getitem__(self, index):
        filename = self.filenames[index]
        img = Image.open(filename).convert('RGB')
        return self.transforms(img)

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


dataset = OleOleDataset(data_dir='./hoge')
dataloader = DataLoader(dataset)

__getitem__() で transforms() に PIL.Image を渡していることからわかるように，torchvisionの前処理は基本的にPIL(Pillow)のラッパになっています． すごく便利ですが，人間は怠惰なので次のような要望が出てきます．

1. PILを使いたくない (≒ NumPy, OpenCVを使いたい)
   • 👼「ndarray に変換すれば使える」
   • 👨「astype('uint8') とか RGB2BGR とかを一々書きたくないです」
2. データローダ外でも使いたい
   • 👨「ネットワークで RGB2GRAY とか幾何的変換してからLossを計算したいのですが」
   • 👼「それはtorchvisionの守備範囲外では」
3. せっかくだからGPU使いたい
   • 👼「cv2.cuda というのがあってだな」
   • 👨「辛そう…」

こんな問答を解決する一手として，今回の主題であります Kornia を以下で紹介します．

Kornia

KorniaはPyTorchをバックエンドとした 微分可能画像処理ライブラリ です． 元々はPyTorch Geometryという名前で幾何学的処理を提供するライブラリでしたが，グラフニューラルネットワークを扱うPyTorch Geometricと名前が競合することや，よりジェネリックな画像処理ライブラリへの拡張を考え，今の名前に改名されたらしいです． OpenCVにインスパイアされている部分が多く，実際に OpenCV.org や Open Source Vision Foundation とパートナーシップを結んでいたりしています*3． 近々ではPyTorch ecosystemの仲間入りを果たしました．

以下では，WACV2020に採択された論文[*4]とライブラリの概観を説明します．なお，ライブラリの構成等は2019/12/08時点のものであり，stableバージョンは 0.1.4.post2 を指します．

kornia.github.io

arxiv.org

Korniaの特徴

論文ではKorniaには次のような特徴があるとしています．

1. 微分可能
2. CPU/GPUで簡単に実行可能
3. ミニバッチでの処理が得意
4. マルチプロセス/マルチノードで実行可能
5. JITコンパイラが優秀

端的に言うと 「PyTorchがバックエンドだから」 という感じです． 殆どの演算がPyTrochの標準モジュール nn.Module を継承しており，一般的なCNNのレイヤと同等に扱うことができるのが一番の強みと思います．

ライブラリ概観

ライブラリは次のようなサブモジュールで構成されています．

サブモジュール	内容
kornia.color	色空間の変換
kornia.filters	画像フィルタリング
kornia.losses	誤差関数及び評価指標
kornia.feature	特徴点検出周りの処理
kornia.geometry	幾何学的変換周りの処理
kornia.utils	ユーティリティ
kornia.contrib	実験的な実装

以下では，各サブモジュールの使い方をザッと見ていきます． テスト画像には lenna を用い，入力サイズはPyTorchの形式に沿って とします．

kornia.color

kornia.color では，主に色空間変換 *5 と輝度やコントラストの調整機能を提供しています． まだstableではないですが，アルファブレンドなども提供予定っぽいです．

import kornia.color as color

# as torch.Tensor (shape: N x 3 x H x W)
org = imread('lenna.png')

# color space conversions
gray = color.rgb_to_grayscale(org)
hsv = color.rgb_to_hsv(org)
hls = color.rgb_to_hls(org)

# color adjustments
brightness = color.adjust_brightness(org, brightness_factor=0.5)
contrast = color.adjust_contrast(org, contrast_factor=0.5)
gamma = color.adjust_gamma(org, gamma=0.5)

結果はこんな感じです．

色空間の変換結果

輝度・コントラスト・ガンマの調整

kornia.filters

kornia.filters では現在ぼかしとエッジ検出用のフィルタリング処理を提供しています． エッジ検出はグレースケールで計算されるため，nn.Sequential() にグレースケール変換とエッジ検出をまとめています． これがまさにnn.Module を継承している強みで，torchvision.transforms.Compose に処理を積む感覚で処理をまとめることができます．

import kornia.color as color
import kornia.filters as filters

# imread as torch.Tensor (shape: N x 3 x H x W)
org = imread('lenna.png')

# blurring
box_blur = filters.box_blur(org, kernel_size=(9, 9))
median_blur = filters.median_blur(org, kernel_size=(9, 9))
gaussian_blur = filters.gaussian_blur2d(org, kernel_size=(9, 9), sigma=(10, 10))

# edge detection
laplacian_detector = nn.Sequential(
    color.RgbToGrayscale(),
    filters.Laplacian(kernel_size=9)
)
laplacian_edge = laplacian_detector(org)

sobel_detector = nn.Sequential(
    color.RgbToGrayscale(),
    filters.Sobel()
)
sobel_edge = sobel_detector(org)

ぼかしフィルタの適用結果

エッジ検出結果

kornia.losses

kornia.losses では公式のnnモジュールに含まれない誤差関数を提供しています． タスクごとにまとめると次のようになります(※印はFuture Releaseです)．

タスク	モジュール名
Object detection	losses.FocalLoss
Semantic segmentation	losses.DiceLoss losses.TverskyLoss
Depth prediction	losses.InverseDepthSmoothnessLoss
画質評価	losses.SSIM losses.TotalVariation※ losses.PSNRLoss※

私個人は画像変換を取り扱うことが多いので，画質評価指標が手軽に使えるのが嬉しいです*6． kornia.losses はそのまま誤差関数として使うようできているのでSSIMやPSNRみたいな類似度も相違度として定義されてます．すなわちこれらの類似度は大小を反転して利用する必要があります．

kornia.feature

kornia.feature では特徴点検出・特徴記述周りの機能を提供します．

特徴点検出ではスケールピラミッドの計算やオリエンテーションの計算などいくつか処理を踏むのですが，Korniaでは feature.ScaleSpaceDetector というラッパが用意されており，各段の処理をモジュールとして渡すことで様々な組み合わせの特徴点検出が可能となります． 例えばHarrisのコーナー検出は次のように書けます．

import kornia.feature as feature

# as torch.Tensor (shape: N x 1 x H x W)
org = imread('lenna.png')

detector = feature.ScaleSpaceDetector(
    num_features=50,
    resp_module=feature.CornerHarris(0.04)
)
lafs, resps = detector(org)

num_features は検出する特徴点の数を表します． 本来の特徴点検出では点の数に制限はありませんが，ミニバッチ化の際に固定サイズでなければならないため，常に num_features 個の特徴点が出力されます*7． 検出した点の位置とスケールを可視化すると次のようになります．

Harrisのコーナー検出の結果

加えてオリエンテーションの計算やアフィン領域推定をしたい場合， detector にそれぞれモジュールを追加します．

detector = kornia.feature.ScaleSpaceDetector(
    num_features=50,
    resp_module=kornia.feature.CornerHarris(0.04),
    aff_module=kornia.feature.LAFAffineShapeEstimator(patch_size=19),
    ori_module=kornia.feature.LAFOrienter(patch_size=19)
)

Harrisのコーナー検出+αの結果

検出した特徴点に対してSIFT特徴を記述するには，次のように書きます． 検出したスケールからパッチをサンプリングしてきて128次元のベクトルに記述します．

descriptor = feature.SIFTDescriptor(32)

patches =  feature.extract_patches_from_pyramid(org, lafs)
B, N, CH, H, W = patches.size()

descs = descriptor(patches.view(B * N, CH, H, W)).view(B, N, -1)
print (descs.shape)  # Out: (N x 50 x 128)

kornia.geometry

kornia.geometryでは，幾何学的変換だったり3次元点への変換だったりを提供します． サブモジュールは次のような感じ．やはりPyTorch Geometryだっただけあって， ここが一番手厚いです．

• kornia.geometry.camera
  • Pinhole Camera
  • Perspective Camera
• kornia.geometry.conversions
• kornia.geometry.depth
• kornia.geometry.linalg
• kornia.geometry.transform
• kornia.geometry.warp

僕が3Dとか何もわからない ライブラリを追いきれないので，今回はホモグラフィ行列による画像変換をやってみます． OpenCVと同様に，変換前後の4点座標からホモグラフィ行列を計算し，画像変換をするには次のようになります．

import kornia.geometry as geometry

org = imread('lenna.png')
n, c, h, w = org .shape

# 変換前後の4点座標
points_src = torch.FloatTensor([[
    [0, 0], [512, 0], [512, 512], [0, 512]]])
points_dst = torch.FloatTensor([[
    [0, 0], [300, 150], [400, 300], [250, 500]]])

M = kornia.get_perspective_transform(points_src, points_dst)
img_warp = kornia.warp_perspective(org, M, dsize=(h, w))

ホモグラフィ行列による変換結果

論文では画像のレジストレーションで著名なLucas–Kanade法を，kornia.geometryとPyTorchの自動微分&勾配降下を利用して実行する例が示されています． 機会があれば同じ要領で画像勾配を用いた画像合成なども試せたらいいかなと思います．

まとめ

本エントリではPyTorchをバックエンドとした微分可能画像処理ライブラリ Kornia を紹介しました． 現状バージョン 0.1.4.post2 と低く発展途上ですが，部分的に使える機能も多い印象です*8． このようなエコシステムが成熟すれば研究の幅も広がるので期待大です．

まぁ結局何が言いたいかというと，

CV勢，コントリビュートしましょう！

*1:抜き去るどころか… - https://preferred.jp/ja/news/pr20191205/

*2:CVPR2019の某ワークショップでは約2/3がPyTorchでした．

*3:というか論文のラストオーサーは完全にOpenCVの中の人

*4:Riba, Edgar, et al. "Kornia: an Open Source Differentiable Computer Vision Library for PyTorch." 2020 IEEE Winter Conference on Applications of Computer Vision (WACV). IEEE, 2020.

*5:密かにRGB2YCbCrのPR投げてました．これ以上は聞かないで．

*6:今までは自分で書いたり，0.3系で書かれたレポジトリをコピーする事がよくありました．

*7:特徴点が検出されないような均質な画像に対しても同様です．

*8:実際に今進めているプロジェクトに使っていたり使っていなかったり…