ノイズや落書きいっぱいの画像や、ズームでぶれた画像を復元するCNN (U-net, DnCNN, WIN5-RB)

はじめに

汚れた画像を綺麗にしたり、落書きを画像から消したり、ズームで品質が落ちた画像を復元したりするために、最近深層学習が使われる方法です。

原理は簡単、ノイズのある画像を変換器に入力して、綺麗な画像を出力させるだけです。

そのような変換器というのはよく畳み込みニューラルネットワーク（CNN）から成なされるものです。

この記事では、そのようなCNNをpytorchで実装してみて、結果を見てみます。

このようなCNNのモデルは色々ありますが、今回は3種のCNNモデルで試して結果を比べてみます。

概要

ノイズのある画像から綺麗な画像に変える変換器を作るためには、まずCNNモデルを決めることです。

そのCNNはノイズのある画像を入力として、出力は入力と同じ大きさの画像。

モデル通りCNNを作ったら、まず最初はCNNの中のパラメータはランダムだから最初の出力は出鱈目であるはず。

その結果は模範の画像と比べてMSE（mean square error）を計算して誤差逆伝播法でパラメータを更新していって、出力は綺麗になっていく。

結局、出力の画像は模範とほぼ変わらないものになります。

学習完了した後、学習に使われなかった画像を入力しても、同じようにノイズが消されるでしょう。（そうじゃない場合は過学習になっているということ）

もっと詳しくはqiitaで参考になる記事がたくさんあります。

• 畳み込みオートエンコーダによる画像の再現、ノイズ除去、セグメンテーション
• ディープラーニングで落書きを消す
• レイトレーシングのデノイズにAIを使う

使う画像データの準備

今回使う画像はこの前DCGANの実装を試したと同じくsafebooruから取得したものです。

詳しくはDCGANの記事で

画像をダウンロードして、256x256サイズに切り取って、このような画像ができます。

そしてこのように3つの方法で画像をいじります。

• ガウス雑音

• 落書き

• 縮小して再び拡大

どちらにせよ、ノイズを加えて画像の品質を落とすことです。そのような画像をCNNで復元してみます。

ただし、CNNの学習をする時に予めこのような画像を予め準備するのではなく、毎回綺麗なサンプルからランダムでノイズを入れて使います。こうやって毎回ランダムで違うノイズからなる画像が使われます。

使うモデル

今回使おうとしたのは4種のモデル

U-net
論文 Ronneberger & Fischer & Brox 2015
コード https://github.com/DuFanXin/U-net , https://github.com/milesial/Pytorch-UNet

RED-Net
論文 Mao & Shen & Yang 2016
コード https://github.com/ved27/RED-net

DnCNN
論文 Zhang et al. 2017
コード https://github.com/cszn/DnCNN

WIN5-RB
論文 Liu & Fang 2017
コード https://github.com/cswin/WIN

ただしRED-Netは結局大きすぎて、使おうとしてみましたが、時間がかかりすぎて効果的ではないため、結局諦めました。

各モデルの構造を絵で説明します。
画像はadobe illustratorのextendscriptで書いてみたものです。

U-net

RED-Net

DnCNN

WIN5-RB

実装

今回実装してみたコード

from glob import glob
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.image import imsave,imread
from skimage.transform import resize
import os,time,torch
from torch import nn
mse = nn.MSELoss()

## ニューラルネットワークモデル ##

class Unet(nn.Module):
    def __init__(self,cn=3):
        super(Unet,self).__init__()

        self.copu1 = nn.Sequential(
            nn.Conv2d(cn,48,3,stride=1,padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(48,48,3,padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2)
        )

        for i in range(2,6):
            self.add_module('copu%d'%i,
                nn.Sequential(
                    nn.Conv2d(48,48,3,stride=1,padding=1),
                    nn.ReLU(inplace=True),
                    nn.MaxPool2d(2)
                )
            )

        self.coasa1 = nn.Sequential(
            nn.Conv2d(48,48,3,stride=1,padding=1),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(48,48,3,stride=2,padding=1,output_padding=1)
        )

        self.coasa2 = nn.Sequential(
            nn.Conv2d(96,96,3,stride=1,padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(96,96,3,stride=1,padding=1),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(96,96,3,stride=2,padding=1,output_padding=1)
        )

        for i in range(3,6):
            self.add_module('coasa%d'%i,
                nn.Sequential(
                    nn.Conv2d(144,96,3,stride=1,padding=1),
                    nn.ReLU(inplace=True),
                    nn.Conv2d(96,96,3,stride=1,padding=1),
                    nn.ReLU(inplace=True),
                    nn.ConvTranspose2d(96,96,3,stride=2,padding=1,output_padding=1)
                )
            )

        self.coli = nn.Sequential(
            nn.Conv2d(96+cn,64,3,stride=1,padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(64,32,3,stride=1,padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(32,cn,3,stride=1,padding=1),
            nn.LeakyReLU(0.1)
        )

        for l in self.modules(): # 重みの初期値
            if(type(l) in (nn.ConvTranspose2d,nn.Conv2d)):
                nn.init.kaiming_normal_(l.weight.data)
                l.bias.data.zero_()

    def forward(self,x):
        x1 = self.copu1(x)
        x2 = self.copu2(x1)
        x3 = self.copu3(x2)
        x4 = self.copu4(x3)
        x5 = self.copu5(x4)

        z = self.coasa1(x5)
        z = self.coasa2(torch.cat((z,x4),1))
        z = self.coasa3(torch.cat((z,x3),1))
        z = self.coasa4(torch.cat((z,x2),1))
        z = self.coasa5(torch.cat((z,x1),1))

        return self.coli(torch.cat((z,x),1))



class DnCNN(nn.Sequential):
    def __init__(self,cn=3):
        super(DnCNN,self).__init__()
        
        self.add_module('coba1',nn.Sequential(
                nn.Conv2d(cn,64,3,stride=1,padding=1),
                nn.ReLU(inplace=True),
            )
        )

        for i in range(2,17):
            self.add_module('coba%d'%i,
                nn.Sequential(
                    nn.Conv2d(64,64,3,stride=1,padding=1),
                    nn.BatchNorm2d(64),
                    nn.ReLU(inplace=True),
                )
            )
        
        self.add_module('coba17',nn.Conv2d(64,cn,3,stride=1,padding=1))
        
        for l in self.modules(): # 重みの初期値
            if(type(l)==nn.Conv2d):
                nn.init.kaiming_normal_(l.weight.data)
                l.bias.data.zero_()
            elif(type(l)==nn.BatchNorm2d):
                l.weight.data.fill_(1)
                l.bias.data.zero_()



class Win5RB(nn.Module):
    def __init__(self,cn=3):
        super(Win5RB,self).__init__()

        inc = cn
        for i in range(1,5):
            self.add_module('coba%d'%i,
                nn.Sequential(
                    nn.Conv2d(inc,64,7,stride=1,padding=3),
                    nn.BatchNorm2d(64),
                    nn.ReLU(inplace=True),
                )
            )
            inc = 64
        
        self.add_module('coba5',nn.Sequential(
                nn.Conv2d(64,cn,7,stride=1,padding=3),
                nn.BatchNorm2d(cn)
            )
        )
        
        for l in self.modules(): # 重みの初期値
            if(type(l)==nn.Conv2d):
                nn.init.kaiming_normal_(l.weight.data)
                l.bias.data.zero_()
            elif(type(l)==nn.BatchNorm2d):
                l.weight.data.fill_(1)
                l.bias.data.zero_()

    def forward(self,x):
        z = self.coba1(x)
        z = self.coba2(z)
        z = self.coba3(z)
        z = self.coba4(z)
        return self.coba5(z) + x



class Rednet(nn.Module):
    def __init__(self,cn=3):
        super(Rednet,self).__init__()

        self.add_module('con1',nn.Sequential(
                nn.Conv2d(cn,64,7,stride=1,padding=3),
                nn.ReLU(inplace=True)
            )
        )

        for i in range(2,16):
            self.add_module('con%d'%i,
                nn.Sequential(
                    nn.Conv2d(64,64,7,stride=1,padding=3),
                    nn.ReLU(inplace=True)
                )
            )

        for i in range(1,15):
            self.add_module('cont%d'%i,
                nn.Sequential(
                    nn.ConvTranspose2d(64,64,7,stride=1,padding=3),
                    nn.ReLU(inplace=True)
                )
            )

        self.add_module('cont15',nn.ConvTranspose2d(64,cn,7,stride=1,padding=3))

        for l in self.modules(): # 重みの初期値
            if(type(l) in (nn.ConvTranspose2d,nn.Conv2d)):
                nn.init.kaiming_normal_(l.weight.data)
                l.bias.data.zero_()

    def forward(self,x):
        z = self.con1(x)
        x2 = self.con2(z)
        z = self.con3(x2)
        x4 = self.con4(z)
        z = self.con5(x4)
        x6 = self.con6(z)
        z = self.con7(x6)
        x8 = self.con8(z)
        z = self.con9(x8)
        x10 = self.con10(z)
        z = self.con11(x10)
        x12 = self.con12(z)
        z = self.con13(x12)
        x14 = self.con14(z)
        z = self.con15(x14)

        z = self.cont1(z)
        z = self.cont2(x14+z)
        z = self.cont3(z)
        z = self.cont4(x12+z)
        z = self.cont5(z)
        z = self.cont6(x10+z)
        z = self.cont7(z)
        z = self.cont8(x8+z)
        z = self.cont9(z)
        z = self.cont10(x6+z)
        z = self.cont11(z)
        z = self.cont12(x4+z)
        z = self.cont13(z)
        z = self.cont14(x2+z)
        z = self.cont15(z)

        return z


## データローダ ##

class Gazoudalo:
    def __init__(self,folder,fnoise,px=256,n_batch=4,random=False):
        self.folder = folder # データの保存したフォルダ
        self.fnoise = fnoise # 使うノイズ関数
        self.px = px # 画像の大きさ
        self.n_batch = n_batch # バッチサイズ
        self.random = random # ランダムするかどうか
        self.file = np.sort(glob(os.path.join(folder,'*.jpg'))) # フォルダの中の全部の画像の名前
        self.len = len(self.file) # 画像の枚数
        self.nkai = int(np.ceil(self.len/n_batch)) # このバッチサイズで繰り返すことができる回数
    
    def __iter__(self):
        self.i_iter = 0
        if(self.random):
            np.random.seed(None)
            # ランダムで順番を入れ替える
            self.i_rand = np.random.permutation(self.len)
        else:
            np.random.seed(0)
            self.i_rand = np.arange(self.len)
        return self
    
    def __next__(self):
        if(self.i_iter>=self.len):
            raise StopIteration
        
        x = []
        y = []
        for i in self.i_rand[self.i_iter:self.i_iter+self.n_batch]:
            yi = imread(self.file[i])/255. # 画像読み込み
            if(self.random and np.random.random()>0.5):
                yi = yi[:,::-1] # ランダムで画像の左右を逆転する
            xi = self.fnoise(yi,self.px) # ノイズを入れる
            
            # 指定のサイズに変える
            xi = resize(xi,(self.px,self.px),anti_aliasing=True,mode='constant')
            yi = resize(yi,(self.px,self.px),anti_aliasing=True,mode='constant')
            xi = torch.Tensor(xi.transpose(2,0,1))
            yi = torch.Tensor(yi.transpose(2,0,1))
            x.append(xi)
            y.append(yi)
            
        x = torch.stack(x)
        y = torch.stack(y)
        self.i_iter += self.n_batch
        return x,y


## 画像変換器 ##

class Dinonet:
    def __init__(self,net,cn,hozon_folder,gakushuuritsu=1e-3,gpu=1):
        self.gakushuuritsu = gakushuuritsu
        self.cn = cn
        self.net = net(cn=cn)
        self.opt = torch.optim.Adam(self.net.parameters(),lr=gakushuuritsu)
        if(gpu):
            # GPUを使う場合
            self.dev = torch.device('cuda')
            self.net.cuda()
        else:
            self.dev = torch.device('cpu')
            
        self.hozon_folder = hozon_folder
        # 保存のフォルダが元々なければ予め作っておく
        if(not os.path.exists(hozon_folder)):
            os.mkdir(hozon_folder)
        
        # 訓練されたパラメータと結果を保存するファイル
        netparam_file = os.path.join(hozon_folder,'netparam.pkl')
        if(os.path.exists(netparam_file)):
            # セーブしておいたデータがすでにある場合は先回から続き
            s = torch.load(netparam_file)
            self.net.load_state_dict(s['w'])
            self.opt.load_state_dict(s['o'])
            self.mounankai = s['n']
            self.sonshitsu = s['l']
        else:
            # 最初から開始
            self.mounankai = 0
            self.sonshitsu = []
    
    def gakushuu(self,dalo_kunren,dalo_kenshou,n_kurikaeshi,n_kaku=5,yokukataru=10):
        print('訓練:%d枚 | 検証%d枚'%(dalo_kunren.len,dalo_kenshou.len))
        t0 = time.time()
        kenshou_data = []
        for da_ken in dalo_kenshou:
            kenshou_data.append(da_ken)
        dalo_kunren.random = True
        print('画像の準備に%.3f分かかった'%((t0-time.time())/60))
        print('==学習開始==')
        
        t0 = time.time()
        # 何回も繰り返して訓練する
        for kaime in range(self.mounankai,self.mounankai+n_kurikaeshi):
            # ミニバッチ開始
            for i_batch,(x,y) in enumerate(dalo_kunren):
                z = self.net(x.to(self.dev))
                sonshitsu = mse(z,y.to(self.dev)) # 訓練データの損失
                self.opt.zero_grad()
                sonshitsu.backward()
                self.opt.step()
                
                # 検証データにテスト
                if((i_batch+1)%int(np.ceil(dalo_kunren.nkai/yokukataru))==0 or i_batch==dalo_kunren.nkai-1):
                    self.net.eval()
                    sonshitsu = []
                    if(n_kaku):
                        gazou = []
                        n_kaita = 0
                    
                    for x,y in kenshou_data:
                        z = self.net(x.to(self.dev))
                        # 検証データの損失
                        sonshitsu.append(mse(z,y.to(self.dev)).item())
                        # 検証データからできた一部の画像を書く
                        if(n_kaita<n_kaku):
                            x = x.numpy().transpose(0,2,3,1) # 入力
                            y = y.numpy().transpose(0,2,3,1) # 模範
                            z = np.clip(z.cpu().detach().numpy(),0,1).transpose(0,2,3,1) # 出力
                            for i,(xi,yi,zi) in enumerate(zip(x,y,z)):
                                # [入力、出力、模範]
                                gazou.append(np.vstack([xi,zi,yi]))
                                n_kaita += 1
                                if(n_kaita>=n_kaku):
                                    break
                    sonshitsu = np.mean(sonshitsu)
                    
                    if(n_kaku):
                        gazou = np.hstack(gazou)
                        imsave(os.path.join(self.hozon_folder,'kekka%03d.jpg'%(kaime+1)),gazou)
                    
                    # 今の状態を出力する
                    print('%d:%d/%d ~ 損失:%.4e %.2f分過ぎた'%(kaime+1,i_batch+1,dalo_kunren.nkai,sonshitsu,(time.time()-t0)/60))
                    self.net.train()
            
            # ミニバッチ一回終了
            self.sonshitsu.append(sonshitsu)
            # パラメータや状態を保存する
            sd = dict(w=self.net.state_dict(),o=self.opt.state_dict(),n=kaime+1,l=self.sonshitsu)
            torch.save(sd,os.path.join(self.hozon_folder,'netparam.pkl'))
            
            # 損失（MSE）の変化を表すグラフを書く
            plt.figure(figsize=[5,4])
            plt.gca(ylabel='MSE')
            ar = np.arange(1,kaime+2)
            plt.plot(ar,self.sonshitsu,'#11aa99')
            plt.tight_layout()
            plt.savefig(os.path.join(self.hozon_folder,'graph.png'))
            plt.close()

    def __call__(self,x,n_batch=8):
        self.net.eval()
        x = torch.Tensor(x)
        y = []
        for i in range(0,len(x),n_batch):
            y.append(self.net(x[i:i+n_batch].to(self.dev)).detach().cpu())
        return torch.cat(y).numpy()


## ノイズの関数 ##

# ガウス雑音
class Gaussnoise:
    def __init__(self,a,clip=True):
        self.a = a # 雑音のサイズ
        self.clip = clip # ノイズを追加した後、[0,1]の範囲内にするかどうか
    
    def __call__(self,y,px):
        x = y + np.random.randn(*y.shape)*self.a
        if(self.clip):
            x = np.clip(x,0,1)
        return x

# 落書きノイズ
class Rakugaki:
    def __init__(self,l0,l1,n):
        self.l0 = l0 # 一番短い線の長さ / 画像のピクセル
        self.l1 = l1 # 一番長い線の長さ / 画像のピクセル
        self.n = n # 線の数
    
    def __call__(self,y,px):
        x = y.copy()
        h = min(y.shape[0],y.shape[1])
        l0 = int(h*self.l0) # 一番短い線のピクセル
        l1 = int(h*self.l1) # 一番長い線のピクセル
        
        kk = np.random.randint(0,h,[self.n,2])
        cc = np.random.randint(0,256,[self.n,3])/255.
        ll = np.random.randint(l0,l1+1,self.n)
        for k,c,l in zip(kk,cc,ll):
            p = np.zeros([l,2],dtype=int)
            p1 = np.random.randint(0,4,l-1)
            p[0] = k
            p2 = p1%2
            p3 = p1>1
            p[1:,0] += (p2==0)*np.where(p3,-1,1)
            p[1:,1] += (p2==1)*np.where(p3,-1,1)
            p = p.cumsum(0)%h
            x[p[:,0],p[:,1]] = c
        return x

# 縮小からできたノイズ
class Shukushou:
    def __init__(self,scale=0.1):
        self.scale = scale
    
    def __call__(self,y,px):
        px = int(px*self.scale)
        return resize(y,(px,px),anti_aliasing=True,mode='constant')



## 実行 ##

kunren_folder = 'kunren' # 訓練データのフォルダ
kenshou_folder = 'kenshou' # 検証データのフォルダ
hozon_folder = 'hozon' # 結果を保存するフォルダ
cn = 3 # チャネル数 (3色データ)
n_batch = 8 # バッチサイズ
px = 256 # 画像の大きさ
n_kurikaeshi = 30 # 何回繰り返すか
n_kaku = 6 # 見るために結果の画像を何枚出力する
yokukataru = 10 # 一回の訓練で何回結果を出力する

# 使うモデルを選ぶ
model = Unet
#net = DnCNN
#net = Win5RB

# ノイズを起こす関数を選ぶ
fnoise = Gaussnoise(0.3)
#fnoise = Rakugaki(0.5,5,25)
#fnoise = Shukushou(0.25)

dalo_kunren = Gazoudalo(kunren_folder,fnoise,px,n_batch) # 訓練データ
dalo_kenshou = Gazoudalo(kenshou_folder,fnoise,px,n_batch) # 検証データ
dino = Dinonet(model,cn,hozon_folder)
# 学習開始
dino.gakushuu(dalo_kunren,dalo_kenshou,n_kurikaeshi,n_kaku,yokukataru)

結果

次はあの3種のモデルで、同じく30回練習した後、検証データでテストした結果を見ます。

ガウス雑音

学習の進歩

出力できた画像。上から「入力、DnCNN、U-net、WIN4-RB、オリジナル」の順

落書き

学習の進歩

出力できた画像

縮小から復元

学習の進歩

出力できた画像

纏め

結果から見ると、どのモデルでもある程度ノイズを除去できるようです。

損失の値から見ると、ガウス雑音と縮小の場合はU-netが一番いいようですが、落書きの場合DnCNNの方が一番いいようです。

その他に、綺麗な画像無しで学習することもできるnoise2noiseという方法があります。これについてこの記事で書いています。