SSIMをPytorchで書いてみた

3787 ワード
PyTorch ComputerVision 画像処理 ssim 深層学習 画像処理 テキストリンク

概要

Stereo MatchingでDisparityの精度を確かめる為にSSIMが良く用いられているのだが、gitで見つけたコードがSSIMの値が0~1に収まっていなかったから、自分で書いてみることにした。

SSMIの計算

paddingするとSSIMの結果が悪くなりそうだったので端っこをclippingした
入力のxとyは比較対象で方はTorch Tensor。

    def SSIM(self, x, y,window_size=3):
        C1 = 0.01 ** 2
        C2 = 0.03 ** 2
        clip_size = (window_size -1)/2

        mu_x = nn.functional.avg_pool2d(x, window_size, 1, padding=0)
        mu_y = nn.functional.avg_pool2d(y, window_size, 1, padding=0)

        x = x[:,:,clip_size:-clip_size,clip_size:-clip_size]
        y = y[:,:,clip_size:-clip_size,clip_size:-clip_size]

        sigma_x = nn.functional.avg_pool2d((x  - mu_x)**2, window_size, 1, padding=0)
        sigma_y = nn.functional.avg_pool2d((y - mu_y)** 2, window_size, 1, padding=0)

        sigma_xy = (
            nn.functional.avg_pool2d((x- mu_x) * (y-mu_y), window_size, 1, padding=0)
        )

        mu_x = mu_x[:,:,clip_size:-clip_size,clip_size:-clip_size]
        mu_y = mu_y[:,:,clip_size:-clip_size,clip_size:-clip_size]

        SSIM_n = (2 * mu_x * mu_y + C1) * (2 * sigma_xy + C2)
        SSIM_d = (mu_x ** 2 + mu_y ** 2 + C1) * (sigma_x + sigma_y + C2)

        SSIM = SSIM_n / SSIM_d

        loss = torch.clamp((1 - SSIM) , 0, 2)
        save_image(loss, 'SSIM_GRAY.png')

        return  torch.mean(loss)

Stereo Matchingの精度を計算する場合

入力の左画像

右の画像を推定したDisparityを使ってワープさせた画像
僕の目には全く同じように見えるけど、SSIMはどうなるのか

SSIMのloss
textureが無い領域が白っぽくなってる。きっとdisparityの推定精度が低かったのだろう。

*白い所がlossが高い。黒い所がlossが0

入力を同じ画像にした場合

もちろんlossは0！

左の画像と右の画像を入力した場合

物体がある場所はきちんとlossがが出てますね。
車のフロントガラスは一部黒くなってるので左右の画像で同じ輝度だという事が分かる。

window sizeを5にした場合

window sizeを7にした場合

window sizeを9にした場合

結論

きっと合ってると思うんですけど、間違ってたら指摘してください！

参考文献

https://github.com/mrharicot/monodepth/blob/b76bee4bd12610b482163871b7ff93e931cb5331/monodepth_model.py#L98
Structural similarity
https://en.wikipedia.org/wiki/Structural_similarity