動画の各フレームからエッジ検出を行う（keras実装）

今回したこと

動画の各フレームにおいてsobelフィルタを用いたエッジ検出を行いました．
普通に各フレームでcv2.Sobelをしろとの声も聞こえてくるかと思いますが， なんとなくkerasのconv3Dで実装したかったのでそのようにしました．
その紹介です．

今回は実装したエッジ検出をMoving MNISTデータセットでテストしました．
時間がない人は下記ノートブックを確認してください．

環境

!python -V
# Python 3.6.9
!pip freeze | grep tensorflow
# tensorflow==2.3.0

kerasで各フレームエッジ検出を実装

こちらのpytorch実装を参考にしています．
https://github.com/JunjH/Revisiting_Single_Depth_Estimation/blob/master/sobel.py

参考元のをpytorch -> keras, conv2D -> conv3Dにしただけです．
conv3Dレイヤーをuse_bias=Falseで定義し，sobelフィルタの重みをset_weightsするだけです．
各フレームでx|y方向のエッジ検出を同時に行うためにConv3D(filters=2, kernel_size=(1,3,3), ...となっています．
任意のフレーム長・画像サイズに対応するため，Input(shape=((None, None, None, 1)))です．
出力は各チャネルにx方向のエッジ，y方向のエッジ画像が含まれています．
実装したもの：

from tensorflow.keras.layers import Input, Conv3D
from tensorflow.keras.models import Model
import numpy as np


class Sobel():
    def __init__(self):
        inputs = Input(shape=((None, None, None, 1)))
        x = Conv3D(filters=2, kernel_size=(1,3,3), padding='valid', use_bias=False)(inputs)
        self.edge_conv = Model(inputs=inputs, outputs=x)

        edge_kx = np.array([[1, 0, -1], 
                            [2, 0, -2], 
                            [1, 0, -1]])
        edge_ky = np.array([[1, 2, 1], 
                            [0, 0, 0], 
                            [-1, -2, -1]])
        edge_k = np.stack((edge_kx, edge_ky))

        edge_k = np.transpose(edge_k, (1,2,0)).reshape(1, 3, 3, 1, 2)
        self.edge_conv.layers[1].set_weights([edge_k])
        # self.edge_conv.trainable = False
        # self.edge_conv.compile(optimizer="adam", loss="mse")
        # self.edge_conv.summary()
        # _________________________________________________________________
        # Layer (type)                 Output Shape                 Param #
        # =================================================================
        # input_35 (InputLayer)        (None, None, None, None, 1)  0      
        # _________________________________________________________________
        # conv3d_38 (Conv3D)           (None, None, None, None, 2)  18     
        # =================================================================
        # Total params: 18
        # Trainable params: 0
        # Non-trainable params: 18
        # _________________________________________________________________
    
    def get_gradient(self, x):
        grad = self.edge_conv.predict(x)
        return grad

使うとき

get_gradient = Sobel().get_gradient
edges = get_gradient(input_data)

Moving MNISTでテストした結果

入力の一部 出力の一部

全コードはgistから確認してください．
sobel_keras · GitHub

終わりに

今回はkerasを使用してエッジ検出を行いました．
普通にやるよりは早いのかな？速度面はまた比較してみます（多分しないけど）．
今回実装した目的は損失関数でエッジ検出を使用するためだったので，次は損失関数の実装に取り掛かります．