爱笑的小姐姐 · 2023年08月07日

无惧暗光!| PE-YOLO: 夜视环境物体检测新突破(附源码实现)

导读

TL;DR: 本文核心内容在于解决在暗光照射条件下进行物体检测的问题。作者建立了一个名为PE-YOLO的暗光物体检测框架,它将金字塔增强网络(PENet)与YOLOv3结合在一起,以改进在暗光环境下的物体检测效果。

PENet的主要思路是利用拉普拉斯金字塔将图像分解成四个不同分辨率的组件。比如吸引笔者注意的是能改进暗光条件下的检测性能。具体地,论文中提出了一个详细的处理模块(DPM),是用于提升图像的细节,它包括两个分支,即上下文分支边缘分支。同时,还有一个低频增强滤波器(LEF),用于捕获低频语义信息并防止高频噪声的干扰

image.png

在训练过程中,PE-YOLO采用了端到端的联合训练方法,并且只使用了普通的检测损失来简化训练过程。通过在低光照物体检测数据集ExDark上的实验,结果证明,与其他暗光检测器和低光增强模型相比,PE-YOLO在 mAP(平均精度)达到 78.0%,FPS(每秒帧数)可以达到 53.6 的条件下,表现出更高的性能,能够适应不同的低光照环境下的物体检测。

动机

这一部分主要向大家介绍与暗光增强和恶劣条件下的目标检测相关的前期研究。这些研究主要集中在改进人类的视觉感知恢复图像细节纠正色彩失真,并为如目标检测这样的高级视觉任务提供高质量的图像。

低光照增强

低光照增强的主要目标是通过恢复图像细节和纠正颜色失真来提高人的视觉感知,并为像目标检测这样的高级视觉任务提供高质量的图像。

image.png
Kind

Kind[1],主要通过具有不同照明级别的配对图像进行训练,而无需 GT 标签。

image.png
MBLLEN

MBLLEN[2]是一个多分支低光照增强网络,该网络在不同的层次上提取特征,并通过多分支融合生成输出图像。

image.png

IAT通过动态查询学习构造端到端 Transformer。在低光照增强模型恢复图像的细节后,检测器的效果得到了提升。

尽管以上方法取得了一定效果,但大多数低光照增强模型都很复杂,对检测器的实时性能有很大影响。

恶劣条件下的目标检测

这一点对于机器人的稳健感知至关重要,随着检测技术的不断发展,目前涌现了一些适应某些恶劣条件的稳健目标检测模型。例如一些通过无监督领域适应将检测器从源领域转移到目标领域,使模型适应恶劣的环境。

image.png
IA-YOLO

IA-YOLO[3]可以自适应地增强每一幅图像以提高检测性能。他们提出了一个用于恶劣天气的可微分图像处理(DIP)模块,并使用了一个小型卷积神经网络(CNN-PP)来调整 DIP 的参数。

关于 IA-YOLO 的强化版本,笔者之前也分享过相关解读,大家可以访问:https://mp.weixin.qq.com/s/qPbxjDuPOFSD2zsWAGmLQw 或者直接关注 CVHub,后台搜索 DIAL-Filters 即可获取超详细的技术解读。

image.png
GDIP-YOLO

GDIP-YOLO[4]是再 IA-YOLO 基础上提出的,作者设计了一种门控机制,允许多个 DIP 并行运行。

image.png
MAET

该方法主要提出了一个用于暗物体检测的多任务自动编码转换MAET[5],探索了照明转换背后的潜在空间。

以上所有相关工作均开源了代码,感兴趣的可以去了解下。

方法

如前所述,由于暗光干扰,暗图像的可见度较差,这严重影响了检测器的性能。为解决这个问题,作者提出了一个金字塔增强网络PENet,并与YOLOv3联合,构建了一个暗物体检测框架PE-YOLOPE-YOLO的框架概览如下所示:

image.png

大家从框架图应该也很容易看出,PE-YOLO就是个两阶段的网络,第一阶段用 PENet 去对原始的暗图像进行增强,得到的增强图像用于第二阶段的输入,后续的网络理论上可以换成任意的目标检测器。所以今天的重点还是再 PENet 上,让我们一探究竟。

PENet

PENet 通过拉普拉斯金字塔将图像分解为不同分辨率的组件。在 PENet 中,提出了两种核心组件:

  • 细节处理模块(DPM
  • 低频增强滤波器(LEF

我们首先对输入图像采用 Gaussian 即高斯滤波器(高斯核的大小为 5×5)进行提取,每次进行高斯金字塔运算后,图像的宽度和高度都减半,这意味着分辨率是原始的 1/4。显然,高斯金字塔的下采样操作是不可逆的。为了在上采样后恢复原始的高分辨率图像,需要丢失的信息,而丢失的信息构成了拉普拉斯金字塔的组件。在重建图像时,我们只需要执行对应的的反向操作即可恢复高分辨率图像。

image.png

如上图所示,我们发现,拉普拉斯金字塔从底部到顶部更加关注全局信息,而反过来更加关注局部细节。这些都是图像下采样过程中丢失的信息,也是 PENet 需要增强的对象。通过细节处理模块(DPM)和低频增强滤波器(LEF)来增强组件,DPM 和 LEF 的操作是并行的。我们即将为大家详细介绍这两个组件。通过分解和重构拉普拉斯金字塔,PENet 可以做到轻量化和有效,这有助于提高后续检测器的性能。

细节处理模块

DPM 被设计成两个分支,一个边缘增强分支和一个上下文分支。其中上下文分支通过捕获远距离依赖获取上下文信息,并全局增强信息。边缘分支使用两个不同方向的Sobel算子计算图像梯度,以获得边缘并增强组件的纹理。

上下文分支:我们在获取远程依赖性之前和之后使用一个残差块来处理特征,而残差学习允许通过跳跃连接传输丰富的低频信息。第一个残差块将特征的通道从 3 改变为 32,第二个残差块将特征的通道从 32 改变为 3。捕获场景中的全局信息已经被证明对于低光增强等低级视觉任务是有益的。

边缘分支:Sobel 算子是一种离散算子,它同时使用高斯滤波器和差分导数。它可以通过计算梯度近似来找到边缘。我们在水平和垂直方向上都使用 Sobel 算子,通过卷积滤波器重新提取边缘信息,并使用残差来增强信息流。

image.png

低频增强滤波器

低频增强滤波器(LEF)主要目的是捕获图像组件中的低频信息,因为这些信息包含了图像的大部分语义信息,对于检测器预测非常重要。

LEF 首先通过卷积层将输入组件转换为不同的尺度。然后使用动态低通滤波器和平均池化捕获和过滤低频信息。特定的池化尺寸(如1×1、2×2、3×3和6×6)会被使用,并结合双线性插值采样,以形成不同规模的低通滤波器。这些滤波器最后通过张量拼接被整合,以恢复原始的图像尺寸。

实验

PE-YOLO 的有效性在 ExDark 数据集上通过多次实验进行了验证。以下是实验部分的主要内容和结果:

image.png

从表中可以看出,直接在 YOLOv3 之前使用低光照增强模型并没有显著提高检测性能。然而,PE-YOLO 在 mAP 上比 MBLLEN 高1.2%,比 Zero-DCE 高1.1%,达到了最佳结果。

image.png

从可视化结果可以发现,尽管 MBLLEN 和 Zero DCE 可以显著提高图像的亮度,但同时也放大了图像中的噪声。相比之下,PE-YOLO 主要捕获低光照图像中物体的潜在信息,同时抑制高频成分中的噪声,因此 PE-YOLO 具有更好的检测性能。

总体来说,这些实验结果表明,PE-YOLO 在低光照条件下的物体检测性能优于其他模型,而且在控制噪声和保留物体信息方面表现出色。

总结

PE-YOLO 是一种新颖的暗环境物体检测框架,该框架整合了金字塔增强网络(PENet)和YOLOv3。为了解决在暗光条件下图像可见性差的问题,该方法使用拉普拉斯金字塔将图像分解为具有不同分辨率的多个组件。然后,利用新提出的详细处理模块(DPM)和低频增强滤波器(LEF)增强这些组件的详细信息和低频信息。PE-YOLO以端到端的方式进行训练,无需额外的损失函数。通过在ExDark数据集上进行的实验,PE-YOLO相较于其他的低光照增强模型和暗光物体检测器,显示出了最佳的性能。

class Lap_Pyramid_Conv(nn.Module):  
    def __init__(self, num_high=3, kernel_size=5, channels=3):  
        super().__init__()  
  
        self.num_high = num_high  
        self.kernel = self.gauss_kernel(kernel_size, channels)  
  
    def gauss_kernel(self, kernel_size, channels):  
        kernel = cv2.getGaussianKernel(kernel_size, 0).dot(  
            cv2.getGaussianKernel(kernel_size, 0).T)  
        kernel = torch.FloatTensor(kernel).unsqueeze(0).repeat(  
            channels, 1, 1, 1)  
        kernel = torch.nn.Parameter(data=kernel, requires_grad=False)  
        return kernel  
  
    def conv_gauss(self, x, kernel):  
        n_channels, _, kw, kh = kernel.shape  
        x = torch.nn.functional.pad(x, (kw // 2, kh // 2, kw // 2, kh // 2),  
                                    mode='reflect')  # replicate    # reflect  
        x = torch.nn.functional.conv2d(x, kernel, groups=n_channels)  
        return x  
  
    def downsample(self, x):  
        return x[:, :, ::2, ::2]  
  
    def pyramid_down(self, x):  
        return self.downsample(self.conv_gauss(x, self.kernel))  
  
    def upsample(self, x):  
        up = torch.zeros((x.size(0), x.size(1), x.size(2) * 2, x.size(3) * 2),  
                         device=x.device)  
        up[:, :, ::2, ::2] = x * 4  
  
        return self.conv_gauss(up, self.kernel)  
  
    def pyramid_decom(self, img):  
        self.kernel = self.kernel.to(img.device)  
        current = img  
        pyr = []  
        for _ in range(self.num_high):  
            down = self.pyramid_down(current)  
            up = self.upsample(down)  
            diff = current - up  
            pyr.append(diff)  
            current = down  
        pyr.append(current)  
        return pyr  
  
    def pyramid_recons(self, pyr):  
        image = pyr[0]  
        for level in pyr[1:]:  
            up = self.upsample(image)  
            image = up + level  
        return image  
  
  
class ResidualBlock(nn.Module):  
    def __init__(self, in_features, out_features):  
        super().__init__()  
        self.conv_x = nn.Conv2d(in_features, out_features, 3, padding=1)  
  
        self.block = nn.Sequential(  
            nn.Conv2d(in_features, in_features, 3, padding=1),  
            nn.LeakyReLU(True),  
            nn.Conv2d(in_features, in_features, 3, padding=1),  
        )  
  
    def forward(self, x):  
        return self.conv_x(x + self.block(x))  
  
  
@BACKBONES.register_module()  
class PENet(nn.Module):  
  
    def __init__(self,  
                 num_high=3,  
                 ch_blocks=32,  
                 up_ksize=1,  
                 high_ch=32,  
                 high_ksize=3,  
                 ch_mask=32,  
                 gauss_kernel=5):  
        super().__init__()  
        self.num_high = num_high  
        self.lap_pyramid = Lap_Pyramid_Conv(num_high, gauss_kernel)  
  
        for i in range(0, self.num_high + 1):  
            self.__setattr__('AE_{}'.format(i), AE(3))  
  
    def forward(self, x):  
        pyrs = self.lap_pyramid.pyramid_decom(img=x)  
  
        trans_pyrs = []  
  
        for i in range(self.num_high + 1):  
            trans_pyr = self.__getattr__('AE_{}'.format(i))(  
                pyrs[-1 - i])  
            trans_pyrs.append(trans_pyr)  
        out = self.lap_pyramid.pyramid_recons(trans_pyrs)  
  
        return out  
  
  
class DPM(nn.Module):  
    def __init__(self, inplanes, planes, act=nn.LeakyReLU(negative_slope=0.2, inplace=True), bias=False):  
        super(DPM, self).__init__()  
  
        self.conv_mask = nn.Conv2d(inplanes, 1, kernel_size=1, bias=bias)  
        self.softmax = nn.Softmax(dim=2)  
        self.sigmoid = nn.Sigmoid()  
  
        self.channel_add_conv = nn.Sequential(  
            nn.Conv2d(inplanes, planes, kernel_size=1, bias=bias),  
            act,  
            nn.Conv2d(planes, inplanes, kernel_size=1, bias=bias)  
        )  
  
    def spatial_pool(self, x):  
        batch, channel, height, width = x.size()  
        input_x = x  
        # [N, C, H * W]  
        input_x = input_x.view(batch, channel, height * width)  
        # [N, 1, C, H * W]  
        input_x = input_x.unsqueeze(1)  
        # [N, 1, H, W]  
        context_mask = self.conv_mask(x)  
        # [N, 1, H * W]  
        context_mask = context_mask.view(batch, 1, height * width)  
        # [N, 1, H * W]  
        context_mask = self.softmax(context_mask)  
        # [N, 1, H * W, 1]  
        context_mask = context_mask.unsqueeze(3)  
        # [N, 1, C, 1]  
        context = torch.matmul(input_x, context_mask)  
        # [N, C, 1, 1]  
        context = context.view(batch, channel, 1, 1)  
  
        return context  
  
    def forward(self, x):  
        # [N, C, 1, 1]  
        context = self.spatial_pool(x)  
  
        # [N, C, 1, 1]  
        channel_add_term = self.channel_add_conv(context)  
        x = x + channel_add_term  
        return x  
  
  
import cv2  
from torchvision import transforms  
  
  
def sobel(img):  
    add_x_total = torch.zeros(img.shape)  
  
    for i in range(img.shape[0]):  
        x = img[i, :, :, :].squeeze(0).cpu().numpy().transpose(1, 2, 0)  
  
        x = x * 255  
        x_x = cv2.Sobel(x, cv2.CV_64F, 1, 0)  
        x_y = cv2.Sobel(x, cv2.CV_64F, 0, 1)  
        add_x = cv2.addWeighted(x_x, 0.5, x_y, 0.5, 0)  
        add_x = transforms.ToTensor()(add_x).unsqueeze(0)  
        add_x_total[i, :, :, :] = add_x  
  
    return add_x_total  
  
  
class AE(nn.Module):  
    def __init__(self, n_feat, reduction=8, bias=False, act=nn.LeakyReLU(negative_slope=0.2, inplace=True), groups=1):  
        super(AE, self).__init__()  
  
        self.n_feat = n_feat  
        self.groups = groups  
        self.reduction = reduction  
        self.agg = nn.Conv2d(6,  
                             3,  
                             1,  
                             stride=1,  
                             padding=0,  
                             bias=False)  
        self.conv_edge = nn.Conv2d(3, 3, kernel_size=1, bias=bias)  
  
        self.res1 = ResidualBlock(3, 32)  
        self.res2 = ResidualBlock(32, 3)  
        self.dpm = nn.Sequential(DPM(32, 32))  
  
        self.conv1 = nn.Conv2d(3, 32, kernel_size=1)  
        self.conv2 = nn.Conv2d(32, 3, kernel_size=1)  
        self.lpm = LowPassModule(32)  
        self.fusion = nn.Conv2d(6, 3, kernel_size=1)  
  
    def forward(self, x):  
        s_x = sobel(x)  
        s_x = self.conv_edge(s_x)  
  
        res = self.res1(x)  
        res = self.dpm(res)  
        res = self.res2(res)  
        out = torch.cat([res, s_x + x], dim=1)  
        out = self.agg(out)  
  
        low_fea = self.conv1(x)  
        low_fea = self.lpm(low_fea)  
        low_fea = self.conv2(low_fea)  
  
        out = torch.cat([out, low_fea], dim=1)  
        out = self.fusion(out)  
  
        return out  
  
  
class LowPassModule(nn.Module):  
    def __init__(self, in_channel, sizes=(1, 2, 3, 6)):  
        super().__init__()  
        self.stages = []  
        self.stages = nn.ModuleList([self._make_stage(size) for size in sizes])  
        self.relu = nn.ReLU()  
        ch = in_channel // 4  
        self.channel_splits = [ch, ch, ch, ch]  
  
    def _make_stage(self, size):  
        prior = nn.AdaptiveAvgPool2d(output_size=(size, size))  
        return nn.Sequential(prior)  
  
    def forward(self, feats):  
        h, w = feats.size(2), feats.size(3)  
        feats = torch.split(feats, self.channel_splits, dim=1)  
        priors = [F.upsample(input=self.stages[i](feats[i]), size=(h, w), mode='bilinear') for i in range(4)]  
        bottle = torch.cat(priors, 1)  
        return self.relu(bottle)  

写在最后

References

[1] Kindling the Darkness

[2] MBLLEN

[3] IA-YOLO

[4] GDIP-YOLO

[5] MAET

作者:派派星
文章来源:CVHub

推荐阅读

更多嵌入式AI干货请关注嵌入式AI专栏。欢迎添加极术小姐姐微信(id:aijishu20)加入技术交流群,请备注研究方向。
推荐阅读
关注数
17204
内容数
1239
嵌入式端AI,包括AI算法在推理框架Tengine,MNN,NCNN,PaddlePaddle及相关芯片上的实现。欢迎加入微信交流群,微信号:aijishu20(备注:嵌入式)
目录
极术微信服务号
关注极术微信号
实时接收点赞提醒和评论通知
安谋科技学堂公众号
关注安谋科技学堂
实时获取安谋科技及 Arm 教学资源
安谋科技招聘公众号
关注安谋科技招聘
实时获取安谋科技中国职位信息