卷积去噪编码器训练自己的数据集并进行缺陷检测(Tensorflow)

卷积去噪自动编码器简介

卷积去噪自动编码器原理其实并不复杂，只是在卷积自动编码器中的输入加入了噪声。这样训练后可以用于获得破损(加入噪声)的输入和纯净的输出之间的映射关系。

上图为卷积自动编码器模型。它包含编码（Encoder）和解码(Decoder)两个部分，在我们加入噪声之后，可以用于做缺陷检测之类。

这里详细的我就不介绍了，大家不懂的可以去看看别的博客。

代码实现

首先是util模块代码如下，这块代码是我室友写的，这是真滴猛，在这里感谢我的室友Dr.Tu。

import numpy as np
import os,math,random
import matplotlib.pyplot as plt

class data_loader:
    def __init__(self, root, batch_size=1, shuffle=False):
        self.root = root
        self.batch_size = batch_size
        self.file_list = os.listdir(self.root)
        if shuffle:
            self.file_list = list(np.array(self.file_list)[random.sample(range(0, len(self.file_list)), len(self.file_list))])
        img = plt.imread(self.root + '/' + self.file_list[0])
        self.shape = (len(self.file_list), img.shape[0], img.shape[1])
        self.flag = 0

    def next_batch(self):
        if self.flag + self.batch_size > self.shape[0]:
            self.file_list = list(np.array(self.file_list)[random.sample(range(0, len(self.file_list)), len(self.file_list))])
            self.flag = 0

        output = np.zeros((self.batch_size, self.shape[1], self.shape[2]))
        temp = 0
        for i in range(self.flag, self.flag + self.batch_size):
            output[temp] = plt.imread(self.root + '/' + self.file_list[i])
            temp = temp + 1

        self.flag += self.batch_size

        return output

def norm(img):
    return (img - 127.5) / 127.5

def denorm(img):
    return (img * 127.5) + 127.5

当然这里的代码是通过minist的手写数据集修改的，本人的训练集是Tilda纺织缺陷数据集，输入大小为512*768，需要的自己去翻qiang下载吧。。。

import tensorflow as tf
import os, util, argparse
import numpy as np
import matplotlib.pyplot as plt
import cv2
from PIL import Image
import scipy.misc

def get_one_image(img_dir):
    image = Image.open(img_dir)
    plt.imshow(image)
    image = image.resize([32, 32])
    image_arr = np.array(image)
    return image_arr

train = True#训练为True，测试为False

parser = argparse.ArgumentParser()
parser.add_argument('--dataset', required=False, default='dataset',  help='添狗必死')
parser.add_argument('--train_subfolder', required=False, default=