Gray ScaleをPytorchで書いてみた

背景

PytorchでGlay Scaleしようと思ったら、Functionがなく、一度PILに変換し、Gray Scaleしてから、Torch Tensorに再変換していたので、計算コスト無駄だなって思ってた。
てかGray Scaleくらいサポートしてくれよ。。。

もしあったらゴメンナサイ

コード

ちなみにGray Scaleは色々と種類があるらしいので、自分の好きなパラメータを使えばいいと思います。

    def getGrayImage(self,rgbImg):
        gray = 0.114*rgbImg[:,0,:,:] + 0.587*rgbImg[:,1,:,:] + 0.299*rgbImg[:,2,:,:]
        gray = torch.unsqueeze(gray,1)
        return gray

結論

PytorchでGrayScaleのFunctionあったら教えてください！

参考文献

グレースケール画像のうんちく
https://qiita.com/yoya/items/96c36b069e74398796f3