基于全卷积网络的彩色显微图像光照不均匀校正

参考：论文英文题目：Correction of uneven illumination in color microscopic image based on fully convolutional network

参考论文链接：https://opg.optica.org/oe/fulltext.cfm?uri=oe-29-18-28503&id=457387

 在做显微图像相关任务是，遇到由于显微光源或者光学系统孔阑造成的图像局部阴影问题，也可称作图像渐晕。为了校正这类图像问题，看了一些相关的论文。传统方法的校正总是不太理想颜色偏移啊之类的问题，某类算法可能就针对一定领域，像暗场的光照校正与明场的光照校正方法和原理就完全不同，一般就是估计光照背景，暗场是减去光照背景，明场就是除光照背景。所以就想试试深度学习的方法，有些模型不是端到端的，处理起来还是比较麻烦，接下来看看这个端到端的模型。

图像校正在深度学习领域本质上还是类似于风格迁移。在训练时将调整好的图像作为参考，给模型输入原图让模型去学习这类校正过程。

论文模型

论文模型总体上来说很简单，论文中说是3部。特征编码，特征解码，细节补充。本质就是unet+带残差结构的FCN。

但论文这里对UNet的具体结构进行了改进：

1.图像在输出后模型后采用双线性插值将原始图像缩放到特定大小

2.下采样没有用池化层，而是选择用步长为2的卷积直接将特征图尺寸减半。标准Unet是卷积，relu，在池化下采样，这种结构对于不均匀光照成像有一些缺陷。例如，它使 标准U-Net缺乏全局颜色信息，这导致生成的图像颜色不一致。减少了网络层数并提高了网 络效率。

3.编码也就是下采样阶段，卷积后得到的feturemap通道数没有增加，标准Unet会增加。

4.解码也就是上采样阶段，没有用标准Unet的反卷积上采样，因为通过反卷积重建的高分辨率特征通常 会有“不均匀重叠”，导致在后来重构的图像中出现高频棋盘格伪影或低频伪影。所以这里采用最近邻插值的方法进行上采样，后卷积，避免图像边界的伪影。NN调整大小卷积（最近邻插值与相同卷积）在防止图像伪影方面取得了最佳效果。最重要的是，NN调整大小卷积在保持瓶颈层传递的照明信息方面表现更佳，这被认为有助于预测图像照明分布。

5.在上采样到与输入图像大小一致的featuremap后，来了与原图来了一个跳过卷积。不同于U-Net，我们提出方法中的跳过连接使得特征解码器 模块中的特征图和特征编码器模块中相应位置的特征图（通过ReLU激活）直接相加。通过添加编码 器网络的局部特征，强迫解码器网络预测更多的特征信息，而不是预测特定的语义像素值。这有利于对输入 图像的光分布的表征。

标准U-Net通过使用复制和裁剪操作，在编码器-解码器模块中完成相应通道特征信息的融合。裁剪操作使得编码器-解码器模块中相应的特征图的大小一致。复制操作是相应特征的连接。

6.细节补充模块，5个卷积层重构最终输出图像，中间加了两个skip connection。

卷积核大小一般为3*3大小，并由Relu函数激活。

在最后输出时，我试了以下，relu可加可不加。

7.损失函数使用的时SSIM结构相似性，还有L1损失函数，结构相似性（SSIM）指数是一 种基于感知的度量[56]。它定义了结构信息（即，在空间上彼此接近的像素高度相关）作为独立于照明和对 比度信息的属性，以反映场景中物体的结构。然后，将图像失真建模为照明、对比度和结构的组合，并采用L1来进一步优化重构图像的亮度和颜色。

所以一个深度学习任务，除了模型的输入输出，还需要好的损失函数，光照不均的重建与图像的内容结构，以及亮度有关，所以采用这两个作为损失函数是没啥问题。

数据集

论文中的数据集是自己的，没有公开，通过作者联系邮箱要也没给，那就只能用github上公开的数据集了。

数据集链接：GitHub - pair-kopti/Shading-correction

内容：分为correct（已校正），original（未校正）

这个数据集据说是通过专业人士调整的，质量确实挺高的。Sensor 2020与Slide images下一个就行，建议下第二个，还有下载时最好用git工具下，直接下太大了，下的不完整。

sensor 2020里大小为5.63GB，分类两个大文件夹correct与original，里面包含40个小文件。

每个小文件中大约为100张彩色图片，图片大小为2304*1719，还是蛮大的。

Slide image 中就小一些，1.24GB，也有小文件，大概10个，每个小文件也是包含大约100张图，图片大小同样为2304*1719。

在数据集文件注释中说Slide image是用于深度学习训练的，看自己个人吧，我感觉都能用。

不过这个数据集还是比较单一的，我也试了一下自己拍的图像，还是校正后差的多，还是得需要自己的数据集进行训练，才能用。

论文代码复现

好了，不多说了，下面直接上代码。

首先呢，模型架构

import torch
import torch.nn as nn
import torch.nn.functional as F


class Encoder(nn.Module):
    def __init__(self, in_channel, out_channel):
        super(Encoder, self).__init__()
        # 一层卷积 一层relu 完成下采样
        self.conv_relu = nn.Sequential(
            # 先卷两次
            nn.Conv2d(in_channel, out_channel, kernel_size=3, stride=2, padding=1),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        x = self.conv_relu(x)
        return x


class Decoder(nn.Module):
    def __init__(self, in_channel, out_channel):
        super(Decoder, self).__init__()
        # 相同卷积
        self.conv_relu = nn.Sequential(
            nn.Conv2d(in_channel, out_channel, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channel, out_channel, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        x = self.conv_relu(x)
        x = F.interpolate(x, scale_factor=2, mode='nearest')
        return x


class DS(nn.Module):
    def __init__(self, in_channel, out_channel):
        super(DS, self).__init__()
        self.conv_relu1 = nn.Sequential(
            nn.Conv2d(in_channel, out_channel, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True)
        )
        self.conv_relu2 = nn.Sequential(
            nn.Conv2d(out_channel, out_channel, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True)
        )
        self.conv_relu3 = nn.Sequential(
            nn.Conv2d(out_channel, out_channel, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        x1 = self.conv_relu1(x)
        x2 = x + self.conv_relu2(x1)
        x3 = x1 + self.conv_relu3(x2)
        return x3


class MICnet(nn.Module):
    def __init__(self):
        super(MICnet, self).__init__()
        self.encode1 = Encoder(3, 64)
        self.encode2 = Encoder(64, 64)
        self.encode3 = Encoder(64, 64)
        self.encode4 = Encoder(64, 64)
        self.encode5 = Encoder(64, 64)

        self.decode1 = Decoder(64, 64)
        self.decode2 = Decoder(64, 64)
        self.decode3 = Decoder(64, 64)
        self.decode4 = Decoder(64, 64)

        self.conv_relu1 = nn.Sequential(
            nn.Conv2d(67, 128, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True)
        )
        self.conv_relu2 = nn.Sequential(
            nn.Conv2d(128, 3, kernel_size=3, stride=1, padding=1),
        )
        self.DS = DS(128,128)

    def forward(self, x):
        resize = F.interpolate(x, size=(96, 96), align_corners=False, mode='bilinear')
        # 下采样编码
        encode1 = self.encode1(resize)
        encode2 = self.encode2(encode1)
        encode3 = self.encode2(encode2)
        encode4 = self.encode2(encode3)
        encode5 = self.encode2(encode4)
        # 谷底上采样次
        x_mid = F.interpolate(encode5, scale_factor=2, mode='nearest')
        add1 = x_mid + encode4
        decode1 = self.decode1(add1)
        add2 = decode1 + encode3
        decode2 = self.decode2(add2)
        add3 = decode2 + encode2
        decode3 = self.decode3(add3)
        add4 = decode3 + encode1
        decode4 = self.decode4(add4)
        # 双线性插值和输入尺寸一致  编码，解码结束
        midblock1 = F.interpolate(decode4, size=(x.shape[2], x.shape[3]), mode='bilinear', align_corners=False)
        midblock = torch.cat((x, midblock1), dim=1)

        # 做一次conv进入细节补充模块
        conv1 = self.conv_relu1(midblock)
        ds = self.DS(conv1)
        result = self.conv_relu2(ds)
        return result

模型大概框架就是这样，DS就是残差那一部分，只不过分开写了，encoder与decoder是Unet的一下层。

2.数据集制作

import os
import glob
import cv2
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
import random
import matplotlib.pyplot as plt

# 数据目录
ci_set = "./Shading-correction/Slide images/Corrected (EMPTY-ZERO method)/"
or_set = "./Shading-correction/Slide images/Original/"

ci_dirs = glob.glob(ci_set + '*')
or_dirs = glob.glob(or_set + '*')

or_img = []
ci_img = []

for subdir in or_dirs:
    for filename in os.listdir(subdir):
        img_path = subdir + "/" + filename
        # print(img_path)
        if '.jpg' in img_path:
            or_img.append(img_path)

for subdir in ci_dirs:
    for filename in os.listdir(subdir):
        img_path = subdir + "/" + filename
        if '.jpg' in img_path:
            # print(img_path)
            ci_img.append(img_path)

# 检查长度
# print(len(or_img))  # 994
# print(len(ci_img))  # 1000

ci_img_new = []
for i in range(len(or_img)):
    or_img_name = or_img[i]
    or_img_name = or_img_name[43:]
    ci_img_name = ci_set + or_img_name
    # print(ci_img_name)
    ci_img_new.append(ci_img_name)
# print(len(or_img))  # 994
# print(len(ci_img_new))  # 1000


new_width = 1600
new_height = 1200
for i in range(len(or_img)):
    or_img_ad = or_img[i]
    ci_img_ad = ci_img_new[i]

    or_img_1 = cv2.imread(or_img_ad)
    or_img_resize = cv2.resize(or_img_1,(new_width,new_height),interpolation=cv2.INTER_LINEAR)

    ci_img_1 = cv2.imread(ci_img_ad)
    ci_img_resize = cv2.resize(ci_img_1, (new_width, new_height), interpolation=cv2.INTER_LINEAR)

    or_img_filename = os.path.join('./data/or_data',f'{i+1}.jpg')
    ci_img_filename = os.path.join('./data/ci_data', f'{i + 1}.jpg')
    cv2.imwrite(or_img_filename, or_img_resize)
    cv2.imwrite(ci_img_filename, ci_img_resize)

这里因为之前的数据分散在各个文件夹里，我觉得太麻烦就全拿出来，放在一个文件夹了，or是原图，ci是校正后的图，并且保存在了data文件下。

import os
import glob
from PIL import Image
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
import random
import matplotlib.pyplot as plt

or_data = "./data/or_data/"
ci_data = "./data/ci_data/"

ci_img_dirs = glob.glob(ci_data + '*')
or_img_dirs = glob.glob(or_data + '*')

# print(len(ci_img_dirs))
# print(len(or_img_dirs))

# 打乱
or_img_new = []
ci_img_new = []

index = list(range(len(or_img_dirs)))
random.shuffle(index)

for i in range(len(or_img_dirs)):
    or_img_new.append(or_img_dirs[index[i]])
    ci_img_new.append(ci_img_dirs[index[i]])

train_transformer = transforms.Compose([
    transforms.ToTensor()
])
test_transformer = transforms.Compose([
    transforms.ToTensor()
])


class ICdataSet(Dataset):
    def __init__(self, or_img1, ic_img1, transformer):
        self.or_img = or_img1
        self.ic_img = ic_img1
        self.transformer = transformer

    def __getitem__(self, index):
        or_img1 = self.or_img[index]
        ic_img1 = self.ic_img[index]

        or_img1_open = Image.open(or_img1)
        or_img_tensor = self.transformer(or_img1_open)

        ic_img1_open = Image.open(ic_img1)
        ic_img_tensor = self.transformer(ic_img1_open)
        return or_img_tensor, ic_img_tensor

    def __len__(self):
        return len(self.or_img)


# 划分数据集
# 按照train 8：test 1：valid 1  共994  所以 800，100，94
end1 = 800
end2 = -94

train_or = or_img_new[:end1]
train_ci = ci_img_new[:end1]
# print(len(train_ci))
test_or = or_img_new[end1 + 1:end1 + 1 + 100]
test_ci = ci_img_new[end1 + 1:end1 + 1 + 100]
# print(len(test_ci))
valid_or = or_img_new[-94:]
valid_ci = ci_img_new[-94:]
# print(len(valid_ci))

train_data = ICdataSet(train_or, train_ci, train_transformer)
test_data = ICdataSet(test_or, test_ci, test_transformer)

dl_train = DataLoader(train_data, batch_size=4, shuffle=True)
dl_test = DataLoader(test_data, batch_size=4, shuffle=True)

or1,ci1 = next(iter(dl_train))
plt.figure(figsize=(12,8))
for i,(img1,img2) in enumerate(zip(or1[:4],ci1[:4])):
    #zip 打包为元组
    img1 =img1.permute(1,2,0).numpy()
    img2 = img2.permute(1, 2, 0).numpy()
    plt.subplot(2,4,i+1)
    plt.imshow(img1)
    plt.subplot(2,4,i+5)
    plt.imshow(img2)
plt.show()

接下来就是将分好的数据制作成数据集，并使用dataloader加载。最后检查了一下可以不写。

可以看到数据是对应的，上面是原图，下面是校正后的图。这样就可以通过dataloder将图像喂给模型。

4.训练函数

import torch
import torch.nn as nn
import torch.nn.functional as F
import os
import glob
from PIL import Image
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
import random
from pytorch_ssim import ssim
from tqdm import tqdm
import torch.optim.lr_scheduler as lrs



# 损失函数计算
def com_loss_fn(y_pred, y, sigma):
    ssim_loss = 1 - ssim(y_pred,y)
    l1_loss = F.l1_loss(y_pred,y)
    batch_loss = sigma * ssim_loss + (1 - sigma) * l1_loss
    return ssim_loss, l1_loss, batch_loss


# 训练函数
def train_epoch(model, trainloader, testloader, device, optimizer, epoch, sigma, scheduler):
    SSIM_loss = 0
    L1_loss = 0
    Loss = 0
    model.train()
    for x, y in tqdm(trainloader):
        x, y = x.to(device), y.to(device)
        y_pred = model(x)
        # 损失函数
        batch_N = y.size(0)
        ssim_loss,l1_loss,loss = com_loss_fn(y_pred, y, sigma)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        with torch.no_grad():
            SSIM_loss += ssim_loss.item()
            L1_loss += l1_loss.item()
            Loss += loss.item

    mean_SSIM_loss = SSIM_loss / len(trainloader.dataset)
    mean_L1_loss = L1_loss / len(trainloader.dataset)
    mean_Loss = Loss / len(trainloader.dataset)
    scheduler.step()

    test_SSIM_loss = 0
    test_L1_loss = 0
    test_Loss = 0

    model.eval()
    with torch.no_grad():
        for x, y in tqdm(testloader):
            x, y = x.to(device), y.to(device)
            y_pred = model(x)
            #损失计算
            batch_N = y.size(0)
            test_ssim_loss, test_l1_loss, test_loss = com_loss_fn(y_pred, y, sigma)

            test_SSIM_loss += test_ssim_loss.item()
            test_L1_loss += test_l1_loss.item()
            test_Loss += test_loss.item()

    test_mean_SSIM_loss = test_SSIM_loss / len(testloader.dataset)
    test_mean_L1_loss = test_L1_loss / len(testloader.dataset)
    test_mean_Loss = test_Loss / len(testloader.dataset)

    staic_dict = model.state_dict()
    torch.save(staic_dict, './checkpoint/{}_train_loss_{}_test_loss_{}.pth'.format(epoch, mean_Loss, test_mean_Loss))
    torch.save(optimizer.state_dict(), './checkpoint/{}the epoch optimizer.pth'.format(epoch))
    print('Learning rate:', scheduler.get_last_lr()[0])
    print(
        'epoch', epoch,
        'train_ssim_loss', mean_SSIM_loss,
        'train_l1_loss', mean_L1_loss,
        'train_loss', mean_Loss,
        'test_ssim_loss', test_mean_SSIM_loss,
        'test_l1_loss', test_mean_L1_loss,
        'test_loss', test_mean_Loss,
    )
    return mean_SSIM_loss, mean_L1_loss, mean_Loss, test_mean_SSIM_loss, test_mean_L1_loss, test_mean_Loss


# 载入数据制作数据集
or_data = "./data/or_data/"
ci_data = "./data/ci_data/"

ci_img_dirs = glob.glob(ci_data + '*')
or_img_dirs = glob.glob(or_data + '*')

# print(len(ci_img_dirs))
# print(len(or_img_dirs))

# 打乱
or_img_new = []
ci_img_new = []

index = list(range(len(or_img_dirs)))
random.shuffle(index)

for i in range(len(or_img_dirs)):
    or_img_new.append(or_img_dirs[index[i]])
    ci_img_new.append(ci_img_dirs[index[i]])


end1 = 800
end2 = -94

train_or = or_img_new[:end1]
train_ci = ci_img_new[:end1]
# print(len(train_ci))
test_or = or_img_new[end1 + 1:end1 + 1 + 100]
test_ci = ci_img_new[end1 + 1:end1 + 1 + 100]
# print(len(test_ci))
valid_or = or_img_new[-94:]
valid_ci = ci_img_new[-94:]
# 将valid信息写进txt文件备用：
valid_or_file = open('valid_or.txt', 'w')
for item in valid_or:
    valid_or_file.write(str(item) + '\n')
valid_or_file.close()

valid_ci_file = open('valid_ci.txt', 'w')
for item in valid_ci:
    valid_ci_file.write(str(item) + '\n')
valid_ci_file.close()

# print(len(valid_ci))

train_transformer = transforms.Compose([
    transforms.ToTensor()
])
test_transformer = transforms.Compose([
    transforms.ToTensor()
])

train_data = ICdataSet(train_or, train_ci, train_transformer)
test_data = ICdataSet(test_or, test_ci, test_transformer)

dl_train = DataLoader(train_data, batch_size=4, shuffle=True)
dl_test = DataLoader(test_data, batch_size=4, shuffle=True)

# 主要参数
model = MICnet()
device = 'cpu'
model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)
sigma = 0.4
scheduler = lrs.StepLR(optimizer, step_size=100, gamma=0.9)

epochs = 120
for epoch in range(epochs):
    train_epoch(model, dl_train, dl_test, device, optimizer, epoch, sigma, scheduler)

这里损失函数ssim这里调了一个包pytorch_ssim。

需要的在这里取链接：https://pan.baidu.com/s/1kirMcwq_zsBjH7acaaXHAw?pwd=gcxy 
提取码：gcxy

还有就是制作数据集时是打乱再分成train，test，vaild，为了防止训练后不知道那个是valid，挣了两个txt，一个存放valid_ci,校正后的图像名地址，还有valid_or原图名和地址，在后面模型验证时调用。

最后模型验证：

import torch
import torch.nn as nn
import torch.nn.functional as F
import os
import glob
from PIL import Image
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
import random
from pytorch_ssim import ssim
from tqdm import tqdm
import torch.optim.lr_scheduler as lrs
from model_train import MICnet,ICdataSet
import matplotlib.pyplot as plt


#从valid_ci与valid_or,读取验证集文件地址
with open('valid_or.txt','r') as valid_or_dir:
    valid_or_dirs = valid_or_dir.readlines()
#去除换行符
valid_or_dirs = [valid_or_dir.strip() for valid_or_dir in valid_or_dirs]
print(valid_or_dirs)

with open('valid_ci.txt','r') as valid_ci_dir:
    valid_ci_dirs = valid_ci_dir.readlines()
#去除换行符
valid_ci_dirs = [valid_ci_dir.strip() for valid_ci_dir in valid_ci_dirs]
print(valid_ci_dirs)

valid_transformer = transforms.Compose([
    transforms.ToTensor()
])

valid_data = ICdataSet(valid_or_dirs,valid_ci_dirs,valid_transformer)

dl_valid = DataLoader(valid_data,batch_size=1,shuffle=False)

or_img,ci_img = next(iter(dl_valid))
or_img = or_img.to('cuda')

model = MICnet()
start_dict = torch.load('')
model = model.to('cuda')
model.eval()

pred_ci_img = model(or_img)
print(pred_ci_img.shape)
pred_ci_img = pred_ci_img.squeeze(0)


plt.figure(figsize=(20, 20))
plt.subplot(1,2,1)
plt.imshow(or_img.permute(1, 2, 0).cpu().numpy())
plt.subplot(1,2,2)
plt.imshow(pred_ci_img.permute(1, 2, 0).cpu().numpy())

我的配置是3090，24G，输入的图像得resize成1000左右大小才能跑，大图跑起来非常慢，还容易爆显卡。所以在训练时resize的图像小一点好跑。

这篇论文最大作用就是应用于图像拼接场景。但是模型泛化能力不够，如果要用还是需要自己做数据集，但图像渐晕校正这一块，就是因为没有好的校正算法才想求助于深度学习的。所以还是得靠发展。

目前有几个好的思路：

将后面的FCN全卷积换成深度可分离卷积，减少参数量，因为这个模型参数量太大了，大图像进模型就跑不了，尤其是显微图像，一般都2000左右了。

试试vit transformer，就是一种编码解码结构。挺感叹的，深度学习发展至今起始很短，尤其是15年后开始蓬勃发展，一年一个大成果，各种模型网络井喷式出现，需要学习的也越来越多，目前检测，分割都已经到头了，大模型一出来解决了这种小模型泛化能力不强的问题，最后归根到底还是堆参数，模型才能进化，具有涌现的功能。

虽然但是，想要入门的同学还是可以大概看看学学，然后找一个好方向是最好的。