自编码AE 实现图片去马赛克 pytorch

谁践踏了优雅 2023-07-16 03:53 8阅读 0赞

去年老早，曾经写过一个AE的实现，不过写的比较墨迹，不够成熟。今天看到了，就重新写一个。

**一.代码**

1.全代码名称展示

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM5OTM4NjY2_size_16_color_FFFFFF_t_70][]

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM5OTM4NjY2_size_16_color_FFFFFF_t_70 1][]

**2.主程序**

**(一).训练阶段**

（1）.dataset.py

import torch
    import os
    import torch.nn as nn
    from torchvision import transforms
    from torch.utils.data import Dataset
    from PIL import Image
    
    class GetData(Dataset):
        def __init__(self,path0,path1): #得到名字list
            super(GetData,self).__init__()
            self.path0 = path0
            self.path1 = path1
            self.name0_list = os.listdir(self.path0)
            self.name1_list = os.listdir(self.path1)
            self.img2data = transforms.Compose([transforms.ToTensor()])
    
        def __len__(self):
            return len(self.name0_list)
    
        def __getitem__(self, index): #按名取图,index对应批次
            self.name0 = self.name0_list[index]
            self.name1 = self.name1_list[index]
            img0 = Image.open(os.path.join(self.path0, self.name0))
            img1 = Image.open(os.path.join(self.path1, self.name1))
            imgdata0 = self.img2data(img0)
            imgdata1 = self.img2data(img1)
    
            return imgdata0, imgdata1

(2).net.py

import torch
    import torch.nn as nn
    
    #卷积
    class ResConv2d(torch.nn.Module):
    
        def __init__(self, in_channels, out_channels):
            super(ResConv2d, self).__init__()
    
            self.sub_net = torch.nn.Sequential(
                torch.nn.Conv2d(in_channels, in_channels, 1, 1, 0), #利用1x1网络
                torch.nn.BatchNorm2d(in_channels),
                torch.nn.PReLU(),
    
                torch.nn.Conv2d(in_channels, in_channels, 3, 1, 1),
                torch.nn.BatchNorm2d(in_channels),
                torch.nn.PReLU(),
    
                torch.nn.Conv2d(in_channels, in_channels, 1, 1, 0),
                torch.nn.BatchNorm2d(in_channels),
                torch.nn.PReLU(),
            )
    
            self.down_net = torch.nn.Sequential(
                torch.nn.Conv2d(in_channels, out_channels, 4, 2, 1),
                torch.nn.BatchNorm2d(out_channels),
                torch.nn.PReLU()
            )
    
        def forward(self, x):
            y = self.sub_net(x)
            return self.down_net(x + y) #加残差
    
    #反卷积
    class ResConvTranspose2d(torch.nn.Module):
    
        def __init__(self, in_channels, out_channels):
            super(ResConvTranspose2d, self).__init__()
    
            self.sub_net = torch.nn.Sequential(
                torch.nn.ConvTranspose2d(in_channels, in_channels, 1, 1, 0),
                torch.nn.BatchNorm2d(in_channels),
                torch.nn.PReLU(),
    
                torch.nn.ConvTranspose2d(in_channels, in_channels, 3, 1, 1),
                torch.nn.BatchNorm2d(in_channels),
                torch.nn.PReLU(),
    
                torch.nn.ConvTranspose2d(in_channels, in_channels, 1, 1, 0),
                torch.nn.BatchNorm2d(in_channels),
                torch.nn.PReLU(),
            )
    
            self.up_net = torch.nn.Sequential(
                torch.nn.ConvTranspose2d(in_channels, out_channels, 4, 2, 1),
                torch.nn.BatchNorm2d(out_channels),
                torch.nn.PReLU(),
            )
    
        def forward(self, x):
            y = self.sub_net(x)
            return self.up_net(x + y)
    
    #编码
    class EncoderNet(torch.nn.Module):
        def __init__(self):
            super(EncoderNet, self).__init__()
    
            self.sub_net = torch.nn.Sequential(
                ResConv2d(3, 64),  # 32
                ResConv2d(64, 128),  # 16
                ResConv2d(128, 256),  # 8
                ResConv2d(256, 512),  # 4
                ResConv2d(512, 1024),  # 2
                ResConv2d(1024, 20)  # 1
            )
    
        def forward(self, x):
            return self.sub_net(x)
    
    #解码
    class DecoderNet(torch.nn.Module):
        def __init__(self):
            super(DecoderNet, self).__init__()
    
            self.decorder = torch.nn.Sequential(
                torch.nn.ConvTranspose2d(20, 1024, 4, 1, 0),
                ResConvTranspose2d(1024, 512),  # 4
                ResConvTranspose2d(512, 256),  # 8
                ResConvTranspose2d(256, 128),  # 16
                torch.nn.ConvTranspose2d(128, 3, 4, 2, 1)  # 64
            )
    
        def forward(self,x):
            return self.decorder(x)

(3).train.py

import torch
    import net
    import dataset
    import torch.nn as nn
    import os
    import shutil
    from torch.utils.data import DataLoader
    from torchvision.utils import save_image
    
    loss_f = nn.MSELoss()
    class MainNet(nn.Module): 
        def __init__(self):
            super(MainNet,self).__init__()
    
            self.encoder = net.EncoderNet()
            self.decoder = net.DecoderNet()
    
        def forward(self,x1):
            y = self.encoder(x1)
            y_ = self.decoder(y)
            return y_
    
        def AELoss(self, y_, x0):
            return loss_f(y_, x0)
    
    #训练
    class Trainer(nn.Module):
        def __init__(self):
            super(Trainer,self).__init__()
    
            self.main_net = MainNet()
            self.main_net.cuda()
    
            '涉及2种损失，自然就会有对应2个优化器做反向传播'
            ae_parameters = []
            ae_parameters.extend(self.main_net.encoder.parameters())
            ae_parameters.extend(self.main_net.decoder.parameters())
            self.opt_ae = torch.optim.Adam(ae_parameters, lr=1e-3)
    
        def train(self):
            for epoch in range(10000):
                if os.path.exists('./param0/encoder.pkl'):
                    self.main_net.encoder.load_state_dict(torch.load('./param0/encoder.pkl'))
                if os.path.exists('./param0/decoder.pkl'):
                    self.main_net.decoder.load_state_dict(torch.load('./param0/decoder.pkl'))
    
                self.dataloader = DataLoader(dataset.GetData(path0=r'C:\Users\87419\Desktop\data\64',
                                 path1=r'C:\Users\87419\Desktop\data\64_dama'), batch_size=128, shuffle=True)
                count = 0
    
                '每个epoch内都是遍历5万张图，即dataloader数。每count一次，即每次循环都是处理batchsize张'
                'dataloader长度 = 总张数/批次数 ：782 = 50000/64。即loader长度等于每个ecpoch的总count数'
                for img0data, img1data in self.dataloader:
    
                    img0data = img0data.cuda()#把输入的数据加cuda，接下来里面的过程数据自然也就以cuda运行
                    img1data = img1data.cuda()
    
                    count += 1
                    # print('/')
                    # print(len(self.dataloader))
    
                    self.main_net.train()#训练模式
                    '每种做梯度更新反向传播，都要重新加载数据！！！'
                    y_ = self.main_net(img1data)
                    # 生成器VAE损失更新
                    aeloss = self.main_net.AELoss(y_, img0data)
                    self.opt_ae.zero_grad()
                    aeloss.backward()
                    self.opt_ae.step()
    
                    if count%25 == 0:
                        self.main_net.eval() #测试模式
                        if os.path.exists('./param0/encoder_tmp.pkl'):
                            shutil.copyfile('./param0/encoder_tmp.pkl', './param0/encoder.pkl')
                        torch.save(self.main_net.encoder.state_dict(), './param0/encoder.pkl')
                        if os.path.exists('./param0/decoder_tmp.pkl'):
                            shutil.copyfile('./param0/decoder_tmp.pkl', './param0/decoder.pkl')
                        torch.save(self.main_net.decoder.state_dict(), './param0/decoder.pkl')
    
                        save_image(img0data[:1],'./result0/{}_{}_0.jpg'.format(epoch, count))#原图
                        save_image(img1data[:1],'./result0/{}_{}_1.jpg'.format(epoch, count))#遮挡图
                        save_image(y_[:1],'./result0/{}_{}_1_0.jpg'.format(epoch, count)) #生成器vae的输出
    
                        print('epoch:',epoch,'|','count:',count,'|','|','aeloss:',aeloss.item()/len(self.dataloader))
    
    if __name__ == '__main__':
        Trainer().train()

**(二).测试阶段**

（1).dataset\_test.py

import torch
    import os
    import torch.nn as nn
    from torchvision import transforms
    from torch.utils.data import Dataset
    from PIL import Image
    
    class GetData(Dataset):
        def __init__(self,path0): #得到名字list
            super(GetData,self).__init__()
            self.path0 = path0
            self.name0_list = os.listdir(self.path0)
            self.img2data = transforms.Compose([transforms.ToTensor()])
    
        def __len__(self):
            return len(self.name0_list)
    
        def __getitem__(self, index): #按名取图,index对应批次
            self.name0 = self.name0_list[index]
            img0 = Image.open(os.path.join(self.path0, self.name0))
            imgdata0 = self.img2data(img0)
    
            return imgdata0

(2).test.py

import torch
    import net
    import dataset_test
    import torch.nn as nn
    import os
    from torch.utils.data import DataLoader
    from torchvision.utils import save_image
    
    class MainNet(nn.Module):
        def __init__(self):
            super(MainNet,self).__init__()
            self.encoder = net.EncoderNet()
            self.decoder = net.DecoderNet()
    
        def forward(self,x1):
            y = self.encoder(x1)
            y_ = self.decoder(y)
            return y_
    
    class Test(nn.Module):
        def __init__(self):
            super(Test,self).__init__()
    
            self.main_net = MainNet()
            self.main_net.cuda()
    
        def test(self):
            if os.path.exists('./param0/encoder.pkl'):
                self.main_net.encoder.load_state_dict(torch.load('./param0/encoder.pkl'))
            if os.path.exists('./param0/decoder.pkl'):
                self.main_net.decoder.load_state_dict(torch.load('./param0/decoder.pkl'))
    
            self.dataloader = DataLoader(dataset_test.GetData(path0=r'C:\Users\87419\Desktop\data\test'))
            count = 0
            self.main_net.eval()  # 测试模式
            for img0data in self.dataloader:
                img0data = img0data.cuda()
                encoded = self.main_net.encoder(img0data)
                decoded = self.main_net.decoder(encoded)
                count += 1
                save_image(decoded, r'C:\Users\87419\Desktop\data\AE_test_result/{}.jpg'.format(count))
    
    if __name__ == '__main__':
        Test().test()

测试效果如下(没仔细训练，只是意思一下)：

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM5OTM4NjY2_size_16_color_FFFFFF_t_70 2][]

[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM5OTM4NjY2_size_16_color_FFFFFF_t_70]: /images/20230528/b3310981e40d42839cab867c93c3df60.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM5OTM4NjY2_size_16_color_FFFFFF_t_70 1]: /images/20230528/1951577924f94a9fa9d57fc64bb4198d.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM5OTM4NjY2_size_16_color_FFFFFF_t_70 2]: /images/20230528/9a734102d9704b9f86960ae82872e946.png