こしあん
2019-10-04

PyTorchでConvolutionフィルターをやる(エッジ検出やアンシャープマスク)

Pocket
LINEで送る


PyTorchでPILのConvolutionフィルター(エッジ検出やアンシャープマスク)をやりたくなったので、どう実装するか考えてみました。

やりたいこと

PIL/PillowのConvolutionフィルター(ImageFilterなど)の処理をPyTorchの畳み込み演算で再現したい

使う画像はこれ。自分が昔撮ってきたものです。「train.jpg」とします。

畳み込みカーネルの値はこちらの記事のものを使っています。だいたいPILの結果と同じになるはずです。

準備

PyTorchのテンソルに変換するために下記の関数を用意しましょう。

from PIL import Image
import numpy as np
import torch
import torch.nn.functional as F
import torchvision

def load_tensor():
    with Image.open("train.jpg") as img:
        array = np.asarray(img, np.float32).transpose([2, 0, 1]) / 255.0
        tensor = torch.as_tensor(np.expand_dims(array, axis=0))  # rank 4
    return tensor

アンシャープマスク

def sharpen_filter():
    kernel = np.array([[-2, -2, -2], [-2, 32, -2], [-2, -2, -2]], np.float32) / 16.0  # convolution filter
    with torch.no_grad():
        # [out_ch, in_ch, .., ..] : channel wiseに計算
        sharpen_k = torch.as_tensor(kernel.reshape(1, 1, 3, 3))

        color = load_tensor()  # color image [1, 3, H, W]
        # channel-wise conv(大事) 3x3 convなのでPadding=1を入れる
        multiband = [F.conv2d(color[:, i:i + 1,:,:], sharpen_k, padding=1) for i in range(3)]
        sharpened_image = torch.cat(multiband, dim=1)
        torchvision.utils.save_image(sharpened_image, "shapren.jpg")

結果

全面の窓ガラスのあたりが少しくっきりしたイメージがあります。ただ、右上の屋根と空の境界線は少しシャギーになりましたね。

コードは、アンシャープマスク用の畳み込みカーネルを用意します。ただ、カラー画像にConvolutionをかけるときは、畳み込みカーネルを(3, 3, 3, 3)というshapeで用意して一括でConv2dを入れると、最後にチャンネルすべての和を取って、値が溢れたりグレースケールのような出力になってしまいます。チャンネル単位でConv2dをとって、最後のConcatしましょう(大事)

エッジ検出

def edge_detection():
    kernel = np.array([[-1, -1, -1], [-1, 8, -1], [-1, -1, -1]], np.float32)  # convolution filter
    with torch.no_grad():
        # [out_ch, in_ch, .., ..] : channel wiseに計算
        edge_k = torch.as_tensor(kernel.reshape(1, 1, 3, 3))

        # エッジ検出はグレースケール化してからやる
        color = load_tensor()  # color image [1, 3, H, W]
        gray_kernel = np.array([0.299, 0.587, 0.114], np.float32).reshape(1, 3, 1, 1)  # color -> gray kernel
        gray_k = torch.as_tensor(gray_kernel)
        gray = torch.sum(color * gray_k, dim=1, keepdim=True)  # grayscale image [1, 1, H, W]

        # エッジ検出
        edge_image = F.conv2d(gray, edge_k, padding=1)
        torchvision.utils.save_image(edge_image, "edge.jpg")

結果:

いい感じにできていますね。

エッジ検出はもともとグレースケールにかけるものなので、チャンネル単位のConv2dは考える必要はありません。グレースケール化の係数やエッジ検出のカーネルは何パターンかあるので、あくまで一例です。

まとめ

PyTorchだろうと(ここでは紹介しなかったけどKerasだろうと)、PILのようなConvolutionフィルターは頑張れば再現できますよ、ということでした。

Related Posts

argparseに直接dictを読み込ませる怪しいやり方... argparseにコマンドライン引数ではなく、ファイルから読み込んだdictをオーバーラップさせる方法を試してみました。本来のargparseの使い方ではない怪しいやり方ですが、JSONやyamlファイルとの連携が可能なので便利ではないかなと思います。 注意 これは本来のargparseの使い...
PCA Color Augmentationを拡張してTensorFlow/Keras向けに実装した... PCA Color AugmentationはAlexNetの論文に示された画像向けのData Augmentationですが、画像用だけではなく、テンソルの固有値分解をすることで構造化データに対しても使えるようにしてみました。これの解説と効果を書きます。 リポジトリ こちら https://...
PyTorchで複数出力があるモデルの出力の型について... 出力が複数あるモデルの訓練というのは少し複雑なモデルだとよく出てきます。PyTorchでは複数出力のモデルの、出力の型はどうなっているでしょうか。それを見ていきます。中間層の値を取りたい場合も使えます。 サンプルコード import torch from torch import nn cl...
TensorFlowの関数で画像にモザイクを書ける方法... TensorFlow2.0の関数を使い、画像にモザイクをかける方法を紹介します。OpenCVやPILでの書き方はいろいろありますが、TensorFlowでどう書くかはまず出てきませんでした。GPUやTPUでのブーストも使えます。 モザイク付与のアルゴリズム いろいろあるとは思いますが、自分が使...
Numpyだけで複数の画像をタイルし1つの画像にまとめる方法... 「torchvision.utilsのmake_gridやテンソルをタイルして保存するのって便利だよね。でも、いちいちこのためにPyTorchのテンソルに変えるのって面倒だよね」ということで同じことをNumpyでも実装してみました。Numpy配列の扱い方を工夫すればいけます。 コード make...
Pocket
Delicious にシェア

Add a Comment

メールアドレスが公開されることはありません。 * が付いている欄は必須項目です