学习资源站

YOLOv11改进-特殊场景检测篇-低照度增强网络PE-YOLO改进主干(改进暗光条件下的物体检测模型,全网独家首发改进)

一、本文介绍

本文给大家带来的改进机制是低照度 图像增强 网络 PE-YOLO 中的PENet,PENet通过 拉普拉斯金字塔 将图像分解成多个 分辨率 组件 ,增强图像细节和低频信息。它包括一个 细节处理模块(DPM) ,用于通过上下文分支和边缘分支增强图像细节,以及一个 低频增强滤波器(LEF) ,以捕获低频语义并减少高频噪声。同时该网络的发布版本并不完善,存在二次创新的机会,后期我会将其网络进行二次创新,增强低照度的检测 性能 同时该网络发布版本存在Bug我也已经修复 欢迎大家订阅本专栏,本专栏每周更新3-5篇最新机制,更有包含我所有改进的文件和交流群提供给大家。

欢迎大家订阅我的专栏一起学习YOLO!


二、PE-YOLO 算法 原理

论文地址: 官方论文地址

代码地址: 官方代码地址


2.1   PE-YOLO的基本原理

PE-YOLO 是一种 改进的暗光条件下的物体检测模型 。它结合了金字塔增强网络(PENet)和 YOLOv3 。PENet通过 拉普拉斯金字塔 将图像分解成多个分辨率的组件,增强图像细节和低频信息。它包括一个 细节处理模块(DPM) ,用于通过上下文分支和边缘分支增强图像细节,以及一个 低频增强滤波器(LEF) ,以捕获低频语义并减少高频噪声。PE-YOLO采用端到端的训练方法,简化训练过程。

PE-YOLO的 基本原理 可以分为几个关键点:

1. 金字塔增强网络(PENet): 使用 拉普拉斯金字塔 将图像分解为不同分辨率的组件,以提升细节和低频信息。
2. 细节处理模块(DPM): 包含上下文分支和边缘分支,专门用于增强图像的细节。
3. 低频增强滤波器(LEF): 用于捕获低频语义信息,同时减少高频噪声。

下面为大家展示了 PE-YOLO系统的总览:

它说明了 如何通过拉普拉斯金字塔 将输入图像分解为不同层级(L0到L3),并通过PENet进行处理,最终提升图像质量以便进行物体检测。图中的细节处理模块(DPM)和低频增强滤波器(LEF)协同工作以增强图像。


2.2 金字塔增强网络

金字塔增强网络(Pyramid Enhancement Network) 是PE-YOLO的关键组成部分,用于增强模型对不同尺度的目标的检测能力。

金字塔增强网络主要包括以下几个 关键要点:

1. 多尺度特征金字塔: 金字塔增强网络使用多个不同尺度的特征金字塔,这些金字塔包含了来自不同层级的特征图。这允许PE-YOLO同时检测不同大小的目标,从小尺寸物体到大尺寸物体都可以有效地检测。

2. 特征融合: 金字塔增强网络通过特征融合的方式将来自不同尺度的特征图进行组合。这有助于提高模型对目标的定位和检测准确性,因为不同尺度的信息被有效地整合在一起。

3. 上采样和下采样: 金字塔增强网络还包括上采样和下采样操作,以进一步调整特征金字塔的尺度。上采样用于增加分辨率,以更好地捕捉小目标的细节信息,而下采样则用于减小分辨率,以更好地捕捉大目标的全局信息。

4. 注意力机制: 金字塔增强网络还引入了注意力机制,以使模型能够集中注意力在最重要的特征上,从而进一步提高检测性能。这有助于减少误检和漏检的情况。

总之,金字塔增强网络是PE-YOLO的关键创新之一,通过多尺度特征金字塔、特征融合、上采样、下采样和注意力机制等技术,提高了PE-YOLO模型在目标检测任务中的性能,使其能够更好地应对不同大小和尺度的目标。


2.3 细节处理模块

细节处理模块(Detail Processing Module,简称DPM) 是PE-YOLO目标检测算法的一个关键组件,旨在增强模型对目标的细节信息的感知和处理能力。DPM的主要任务是通过 上下文分支和边缘分支 来对目标进行更详细的处理。

我为大家总结了PE-YOLO中细节处理模块(DPM)的 主要特点和功能:

1. 上下文分支(Context Branch): 上下文分支负责获取上下文信息,通过捕捉远程依赖关系来理解目标周围的环境。这有助于模型更好地理解目标与其周围环境的关系,从而提高目标检测的准确性。上下文信息的引入可以使模型更好地分辨目标和背景之间的区别。

2. 边缘分支(Edge Branch): 边缘分支使用两个Sobel算子(Sobel operators)在不同方向上计算图像的梯度,从而获得目标的边缘信息。这有助于模型更好地识别目标的轮廓和边缘特征,并增强目标组件的纹理信息。边缘信息对于目标的细节识别和检测非常重要。

3. 组件增强: DPM的综合作用是增强目标的各个组件,包括上下文信息的增强和边缘信息的增强。这使得模型更能够准确地捕捉目标的细节特征,从而提高目标检测性能。

下图展示的是DPM的结构包括 上下文分支(CB)和边缘分支(EB):

上下文分支 通过捕捉远程依赖关系来获取上下文信息,并全局增强组件。

边缘分支 使用两个不同方向的Sobel运算符来计算图像梯度,以获取边缘并增强组件的纹理。


2.4 低频增强滤波器

低频增强滤波器(Low-Frequency Enhancement Filter,简称LEF) 用于捕捉和增强图像中的低频信息,这些低频信息通常包含了图像的大部分语义和关键信息,对于检测器的预测非常重要。

PE-YOLO中低频增强滤波器(LEF)的 主要特点和功能 总结如下:

1. 自适应平均池化: LEF使用不同尺寸的自适应平均池化操作来截取低频分量。这意味着LEF可以动态地适应不同尺度和语义的低频信息,以确保最大程度地捕捉图像中的关键细节。

2. 低频信息捕捉: LEF的主要任务是捕捉和增强图像中的低频信息,这些信息包含了图像的主要语义和关键细节。通过使用低通滤波器来过滤特征,LEF只允许低于截止频率的信息通过,从而增强了低频成分。

3. 多尺度处理: 考虑到Inception的多尺度结构,LEF在不同的尺寸上应用自适应平均池化,以适应不同语义和尺度的低频信息。这有助于提高模型对图像细节的理解和捕捉。

4. 通道分离: LEF将特征f分为四个部分,即{f1, f2, f3, f4},通过通道分离的方式,每个部分都可以独立处理,以进一步增强低频信息。

下图展示了 低频增强滤波器(LEF)的详细信息 。LEF由不同大小的自适应平均池化组成,用于截取低频分量。

考虑到Inception的多尺度结构,我们使用了大小分别为1×1、2×2、3×3、6×6的自适应平均池化,并在每个尺度的末尾使用上采样来恢复特征的原始大小。不同核大小的平均池化形成了一个低通滤波器。我们通过通道分离将f分成四个部分,即{f1, f2, f3, f4}。每个部分都使用不同尺寸的池化进行处理,描述如下:

其中 f_{i} 是分割在通道上的f的一部分,Up是双线性插值采样, \beta _{s} 是不同尺寸s × s的自适应平均池化。最后,在张量拼接每个{ f_{i} , i = 1, 2, 3, 4}之后,我们将它们还原为f ∈ R^{h×w×3}


三、PE-YOLO的核心代码

这个代码的使用方式比较特殊,大家需要注意看后面的使用教程!

  1. import torch
  2. import torch.nn as nn
  3. import torch.nn.functional as F
  4. class Lap_Pyramid_Conv(nn.Module):
  5. def __init__(self, num_high=3, kernel_size=5, channels=3):
  6. super().__init__()
  7. self.num_high = num_high
  8. self.kernel = self.gauss_kernel(kernel_size, channels)
  9. def gauss_kernel(self, kernel_size, channels):
  10. kernel = cv2.getGaussianKernel(kernel_size, 0).dot(
  11. cv2.getGaussianKernel(kernel_size, 0).T)
  12. kernel = torch.FloatTensor(kernel).unsqueeze(0).repeat(
  13. channels, 1, 1, 1)
  14. kernel = torch.nn.Parameter(data=kernel, requires_grad=False)
  15. return kernel
  16. def conv_gauss(self, x, kernel):
  17. n_channels, _, kw, kh = kernel.shape
  18. x = torch.nn.functional.pad(x, (kw // 2, kh // 2, kw // 2, kh // 2),
  19. mode='reflect') # replicate # reflect
  20. x = torch.nn.functional.conv2d(x, kernel, groups=n_channels)
  21. return x
  22. def downsample(self, x):
  23. return x[:, :, ::2, ::2]
  24. def pyramid_down(self, x):
  25. return self.downsample(self.conv_gauss(x, self.kernel))
  26. def upsample(self, x):
  27. up = torch.zeros((x.size(0), x.size(1), x.size(2) * 2, x.size(3) * 2),
  28. device=x.device)
  29. up[:, :, ::2, ::2] = x * 4
  30. return self.conv_gauss(up, self.kernel)
  31. def pyramid_decom(self, img):
  32. self.kernel = self.kernel.to(img.device)
  33. current = img
  34. pyr = []
  35. for _ in range(self.num_high):
  36. down = self.pyramid_down(current)
  37. up = self.upsample(down)
  38. diff = current - up
  39. pyr.append(diff)
  40. current = down
  41. pyr.append(current)
  42. return pyr
  43. def pyramid_recons(self, pyr):
  44. image = pyr[0]
  45. for level in pyr[1:]:
  46. up = self.upsample(image)
  47. image = up + level
  48. return image
  49. class ResidualBlock(nn.Module):
  50. def __init__(self, in_features, out_features):
  51. super().__init__()
  52. self.conv_x = nn.Conv2d(in_features, out_features, 3, padding=1)
  53. self.block = nn.Sequential(
  54. nn.Conv2d(in_features, in_features, 3, padding=1),
  55. nn.LeakyReLU(True),
  56. nn.Conv2d(in_features, in_features, 3, padding=1),
  57. )
  58. def forward(self, x):
  59. return self.conv_x(x + self.block(x))
  60. class PENet(nn.Module):
  61. def __init__(self,
  62. num_high=3,
  63. gauss_kernel=5):
  64. super().__init__()
  65. self.num_high = num_high
  66. self.lap_pyramid = Lap_Pyramid_Conv(num_high, gauss_kernel)
  67. for i in range(0, self.num_high + 1):
  68. self.__setattr__('AE_{}'.format(i), AE(3))
  69. def forward(self, x):
  70. pyrs = self.lap_pyramid.pyramid_decom(img=x)
  71. trans_pyrs = []
  72. for i in range(self.num_high + 1):
  73. trans_pyr = self.__getattr__('AE_{}'.format(i))(
  74. pyrs[-1 - i])
  75. trans_pyrs.append(trans_pyr)
  76. out = self.lap_pyramid.pyramid_recons(trans_pyrs)
  77. return out
  78. class DPM(nn.Module):
  79. def __init__(self, inplanes, planes, act=nn.LeakyReLU(negative_slope=0.2, inplace=True), bias=False):
  80. super(DPM, self).__init__()
  81. self.conv_mask = nn.Conv2d(inplanes, 1, kernel_size=1, bias=bias)
  82. self.softmax = nn.Softmax(dim=2)
  83. self.sigmoid = nn.Sigmoid()
  84. self.channel_add_conv = nn.Sequential(
  85. nn.Conv2d(inplanes, planes, kernel_size=1, bias=bias),
  86. act,
  87. nn.Conv2d(planes, inplanes, kernel_size=1, bias=bias)
  88. )
  89. def spatial_pool(self, x):
  90. batch, channel, height, width = x.size()
  91. input_x = x
  92. # [N, C, H * W]
  93. input_x = input_x.view(batch, channel, height * width)
  94. # [N, 1, C, H * W]
  95. input_x = input_x.unsqueeze(1)
  96. # [N, 1, H, W]
  97. context_mask = self.conv_mask(x)
  98. # [N, 1, H * W]
  99. context_mask = context_mask.view(batch, 1, height * width)
  100. # [N, 1, H * W]
  101. context_mask = self.softmax(context_mask)
  102. # [N, 1, H * W, 1]
  103. context_mask = context_mask.unsqueeze(3)
  104. # [N, 1, C, 1]
  105. context = torch.matmul(input_x, context_mask)
  106. # [N, C, 1, 1]
  107. context = context.view(batch, channel, 1, 1)
  108. return context
  109. def forward(self, x):
  110. # [N, C, 1, 1]
  111. context = self.spatial_pool(x)
  112. # [N, C, 1, 1]
  113. channel_add_term = self.channel_add_conv(context)
  114. x = x + channel_add_term
  115. return x
  116. import cv2
  117. from torchvision import transforms
  118. def sobel(img):
  119. device = img.device
  120. add_x_total = torch.zeros(img.shape)
  121. for i in range(img.shape[0]):
  122. x = img[i, :, :, :].squeeze(0).cpu().detach().numpy().transpose(1, 2, 0)
  123. x = x * 255
  124. x_x = cv2.Sobel(x, cv2.CV_64F, 1, 0)
  125. x_y = cv2.Sobel(x, cv2.CV_64F, 0, 1)
  126. add_x = cv2.addWeighted(x_x, 0.5, x_y, 0.5, 0)
  127. add_x = transforms.ToTensor()(add_x).unsqueeze(0)
  128. add_x_total[i, :, :, :] = add_x
  129. add_x_total = add_x_total.to(device)
  130. return add_x_total
  131. class AE(nn.Module):
  132. def __init__(self, n_feat, reduction=8, bias=False, act=nn.LeakyReLU(negative_slope=0.2, inplace=True), groups=1):
  133. super(AE, self).__init__()
  134. self.n_feat = n_feat
  135. self.groups = groups
  136. self.reduction = reduction
  137. self.agg = nn.Conv2d(6,
  138. 3,
  139. 1,
  140. stride=1,
  141. padding=0,
  142. bias=False)
  143. self.conv_edge = nn.Conv2d(3, 3, kernel_size=1, bias=bias)
  144. self.res1 = ResidualBlock(3, 32)
  145. self.res2 = ResidualBlock(32, 3)
  146. self.dpm = nn.Sequential(DPM(32, 32))
  147. self.conv1 = nn.Conv2d(3, 32, kernel_size=1)
  148. self.conv2 = nn.Conv2d(32, 3, kernel_size=1)
  149. self.lpm = LowPassModule(32)
  150. self.fusion = nn.Conv2d(6, 3, kernel_size=1)
  151. def forward(self, x):
  152. s_x = sobel(x)
  153. s_x = self.conv_edge(s_x)
  154. res = self.res1(x)
  155. res = self.dpm(res)
  156. res = self.res2(res)
  157. out = torch.cat([res, s_x + x], dim=1)
  158. out = self.agg(out)
  159. low_fea = self.conv1(x)
  160. low_fea = self.lpm(low_fea)
  161. low_fea = self.conv2(low_fea)
  162. out = torch.cat([out, low_fea], dim=1)
  163. out = self.fusion(out)
  164. return out
  165. class LowPassModule(nn.Module):
  166. def __init__(self, in_channel, sizes=(1, 2, 3, 6)):
  167. super().__init__()
  168. self.stages = []
  169. self.stages = nn.ModuleList([self._make_stage(size) for size in sizes])
  170. self.relu = nn.ReLU()
  171. ch = in_channel // 4
  172. self.channel_splits = [ch, ch, ch, ch]
  173. def _make_stage(self, size):
  174. prior = nn.AdaptiveAvgPool2d(output_size=(size, size))
  175. return nn.Sequential(prior)
  176. def forward(self, feats):
  177. h, w = feats.size(2), feats.size(3)
  178. feats = torch.split(feats, self.channel_splits, dim=1)
  179. priors = [F.upsample(input=self.stages[i](feats[i]), size=(h, w), mode='bilinear') for i in range(4)]
  180. bottle = torch.cat(priors, 1)
  181. return self.relu(bottle)
  182. if __name__ == "__main__":
  183. # Generating Sample image
  184. image_size = (1, 3, 224, 224)
  185. image = torch.rand(*image_size)
  186. # Model
  187. mobilenet_v1 = PENet()
  188. out = mobilenet_v1(image)
  189. print(out.size())

四、PE-YOLO的添加方式

这个添加方式和之前的变了一下,以后的添加方法都按照这个来了,是为了和群内的文件适配。


4.1 修改一

第一还是建立文件,我们找到如下 ultralytics /nn/modules文件夹下建立一个目录名字呢就是'Addmodules'文件夹( 用群内的文件的话已经有了无需新建) !然后在其内部建立一个新的py文件将核心代码复制粘贴进去即可。


4.2 修改二

第二步我们在该目录下创建一个新的py文件名字为'__init__.py'( 用群内的文件的话已经有了无需新建) ,然后在其内部导入我们的检测头如下图所示。


4.3 修改三

第三步我门中到如下文件'ultralytics/nn/tasks.py'进行导入和注册我们的模块( 用群内的文件的话已经有了无需重新导入直接开始第四步即可)

从今天开始以后的教程就都统一成这个样子了,因为我默认大家用了我群内的文件来进行修改!!

到此就完事了注册的工作,该模型无需添加任何参数是一种无参的机制,所以导入进来即可,该网络本身存在一个设备定义的错误,我解决以后需要关闭混合精度验证这是必须的否则会报错,其次该网络修改完打印不出来参数,都可以通过下面的步骤来解决。大家注意看!!


关闭混合精度验证!

找到'ultralytics/engine/validator.py'文件找到 'class BaseValidator:' 然后在其'__call__'中 self.args.half = self.device.type != 'cpu' # force FP16 val during training的一行代码下面加上self.args.half = False

打印计算量的问题!

计算的GFLOPs计算 异常 不打印,所以需要额外修改一处, 我们找到如下文件'ultralytics/utils/torch_utils.py'文件内有如下的代码按照如下的图片进行修改,有一个get_flops的函数我们直接用我给的代码全部替换!

  1. def get_flops(model, imgsz=640):
  2. """Return a YOLO model's FLOPs."""
  3. if not thop:
  4. return 0.0 # if not installed return 0.0 GFLOPs
  5. try:
  6. model = de_parallel(model)
  7. p = next(model.parameters())
  8. if not isinstance(imgsz, list):
  9. imgsz = [imgsz, imgsz] # expand if int/float
  10. try:
  11. # Use stride size for input tensor
  12. stride = 640
  13. im = torch.empty((1, 3, stride, stride), device=p.device) # input image in BCHW format
  14. flops = thop.profile(deepcopy(model), inputs=[im], verbose=False)[0] / 1e9 * 2 # stride GFLOPs
  15. return flops * imgsz[0] / stride * imgsz[1] / stride # imgsz GFLOPs
  16. except Exception:
  17. # Use actual image size for input tensor (i.e. required for RTDETR models)
  18. im = torch.empty((1, p.shape[1], *imgsz), device=p.device) # input image in BCHW format
  19. return thop.profile(deepcopy(model), inputs=[im], verbose=False)[0] / 1e9 * 2 # imgsz GFLOPs
  20. except Exception:
  21. return 0.0


五、PE-YOLO的yaml文件和运行记录

5.1 PE-YOLO的yaml文件

训练信息:YOLO11-PEYOLO summary: 470 layers, 2,682,018 parameters, 2,681,927 gradients, 28.9 GFLOPs

此模型训练需要关闭amp训练,否则会报错.

  1. # Ultralytics YOLO 🚀, AGPL-3.0 license
  2. # YOLO11 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect
  3. # Parameters
  4. nc: 80 # number of classes
  5. scales: # model compound scaling constants, i.e. 'model=yolo11n.yaml' will call yolo11.yaml with scale 'n'
  6. # [depth, width, max_channels]
  7. n: [0.50, 0.25, 1024] # summary: 319 layers, 2624080 parameters, 2624064 gradients, 6.6 GFLOPs
  8. s: [0.50, 0.50, 1024] # summary: 319 layers, 9458752 parameters, 9458736 gradients, 21.7 GFLOPs
  9. m: [0.50, 1.00, 512] # summary: 409 layers, 20114688 parameters, 20114672 gradients, 68.5 GFLOPs
  10. l: [1.00, 1.00, 512] # summary: 631 layers, 25372160 parameters, 25372144 gradients, 87.6 GFLOPs
  11. x: [1.00, 1.50, 512] # summary: 631 layers, 56966176 parameters, 56966160 gradients, 196.0 GFLOPs
  12. # YOLO11n backbone
  13. backbone:
  14. # [from, repeats, module, args]
  15. - [-1, 1, PENet, []] # 0-P1/2
  16. - [-1, 1, Conv, [64, 3, 2]] # 1-P1/2
  17. - [-1, 1, Conv, [128, 3, 2]] # 2-P2/4
  18. - [-1, 2, C3k2, [256, False, 0.25]]
  19. - [-1, 1, Conv, [256, 3, 2]] # 4-P3/8
  20. - [-1, 2, C3k2, [512, False, 0.25]]
  21. - [-1, 1, Conv, [512, 3, 2]] # 6-P4/16
  22. - [-1, 2, C3k2, [512, True]]
  23. - [-1, 1, Conv, [1024, 3, 2]] # 8-P5/32
  24. - [-1, 2, C3k2, [1024, True]]
  25. - [-1, 1, SPPF, [1024, 5]] # 10
  26. - [-1, 2, C2PSA, [1024]] # 11
  27. # YOLO11n head
  28. head:
  29. - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  30. - [[-1, 7], 1, Concat, [1]] # cat backbone P4
  31. - [-1, 2, C3k2, [512, False]] # 14
  32. - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  33. - [[-1, 5], 1, Concat, [1]] # cat backbone P3
  34. - [-1, 2, C3k2, [256, False]] # 17 (P3/8-small)
  35. - [-1, 1, Conv, [256, 3, 2]]
  36. - [[-1, 14], 1, Concat, [1]] # cat head P4
  37. - [-1, 2, C3k2, [512, False]] # 20 (P4/16-medium)
  38. - [-1, 1, Conv, [512, 3, 2]]
  39. - [[-1, 11], 1, Concat, [1]] # cat head P5
  40. - [-1, 2, C3k2, [1024, True]] # 23 (P5/32-large)
  41. - [[17, 20, 23], 1, Detect, [nc]] # Detect(P3, P4, P5)

5.2 PE-YOLO的训练过程截图


五、本文总结

到此本文的正式分享内容就结束了,在这里给大家推荐我的YOLOv11改进有效涨点专栏,本专栏目前为新开的平均质量分98分,后期我会根据各种最新的前沿顶会进行论文复现,也会对一些老的改进机制进行补充,如果大家觉得本文帮助到你了,订阅本专栏,关注后续更多的更新~