【原创创新点】金属工件缺陷检测系统:Efficient Multi-Scale-Conv的改进YOLOv8

1.研究背景与意义

项目参考AAAI Association for the Advancement of Artificial Intelligence

研究背景与意义:金属工件是现代工业生产中不可或缺的重要组成部分。金属工件的质量和性能直接影响到产品的品质和效率,因此对金属工件的研究和改进具有重要的意义。随着科技的不断进步和工业的发展,对金属工件的要求也越来越高,传统的金属加工方法已经无法满足现代工业的需求。因此,寻找原创创新点来改进金属工件的制造和加工技术,提高金属工件的质量和性能,具有重要的理论和实践意义。

金属工件的质量和性能直接关系到产品的品质和效率。传统的金属加工方法存在一些问题,如加工精度不高、工艺复杂、成本高等。这些问题导致金属工件的质量和性能无法达到理想的水平,影响到产品的品质和效率。因此,寻找原创创新点来改进金属工件的制造和加工技术,提高金属工件的质量和性能,对于提高产品的品质和效率具有重要的意义。

金属工件的质量和性能也直接关系到工业生产的效益和竞争力。随着全球经济的发展和竞争的加剧,工业生产的效益和竞争力成为企业追求的目标。而金属工件作为工业生产的重要组成部分,其质量和性能的提高将直接影响到工业生产的效益和竞争力。因此,寻找原创创新点来改进金属工件的制造和加工技术,提高金属工件的质量和性能,对于提高工业生产的效益和竞争力具有重要的意义。

金属工件的质量和性能的提高还能够推动整个工业领域的发展。金属工件广泛应用于汽车、航空航天、机械制造等领域,其质量和性能的提高将直接推动这些领域的发展。例如,提高金属工件的强度和耐磨性可以延长汽车和机械设备的使用寿命,提高航空航天器的安全性能。因此,寻找原创创新点来改进金属工件的制造和加工技术,提高金属工件的质量和性能,对于推动整个工业领域的发展具有重要的意义。

金属工件的质量和性能对于产品的品质和效率、工业生产的效益和竞争力以及整个工业领域的发展都具有重要的影响。因此,寻找原创创新点来改进金属工件的制造和加工技术,提高金属工件的质量和性能,具有重要的理论和实践意义。这不仅可以满足现代工业对金属工件的高要求,还可以推动整个工业领域的发展,促进经济的繁荣和社会的进步。

2.图片演示



3.视频演示

【原创创新点】金属工件缺陷检测系统:Efficient Multi-Scale-Conv的改进YOLOv8_哔哩哔哩_bilibili

4.数据集的采集&标注和整理

图片的收集

首先,我们需要收集所需的图片。这可以通过不同的方式来实现,例如使用现有的公开数据集MDDatasets。

【原创创新点】金属工件缺陷检测系统:Efficient Multi-Scale-Conv的改进YOLOv8_第1张图片

labelImg是一个图形化的图像注释工具,支持VOC和YOLO格式。以下是使用labelImg将图片标注为VOC格式的步骤:

(1)下载并安装labelImg。
(2)打开labelImg并选择“Open Dir”来选择你的图片目录。
(3)为你的目标对象设置标签名称。
(4)在图片上绘制矩形框,选择对应的标签。
(5)保存标注信息,这将在图片目录下生成一个与图片同名的XML文件。
(6)重复此过程,直到所有的图片都标注完毕。

由于YOLO使用的是txt格式的标注,我们需要将VOC格式转换为YOLO格式。可以使用各种转换工具或脚本来实现。

下面是一个简单的方法是使用Python脚本,该脚本读取XML文件,然后将其转换为YOLO所需的txt格式。

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import xml.etree.ElementTree as ET
import os

classes = []  # 初始化为空列表

CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))

def convert(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)

def convert_annotation(image_id):
    in_file = open('./label_xml\%s.xml' % (image_id), encoding='UTF-8')
    out_file = open('./label_txt\%s.txt' % (image_id), 'w')  # 生成txt格式文件
    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)

    for obj in root.iter('object'):
        cls = obj.find('name').text
        if cls not in classes:
            classes.append(cls)  # 如果类别不存在,添加到classes列表中
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
             float(xmlbox.find('ymax').text))
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')

xml_path = os.path.join(CURRENT_DIR, './label_xml/')

# xml list
img_xmls = os.listdir(xml_path)
for img_xml in img_xmls:
    label_name = img_xml.split('.')[0]
    print(label_name)
    convert_annotation(label_name)

print("Classes:")  # 打印最终的classes列表
print(classes)  # 打印最终的classes列表

整理数据文件夹结构

我们需要将数据集整理为以下结构:

-----data
   |-----train
   |   |-----images
   |   |-----labels
   |
   |-----valid
   |   |-----images
   |   |-----labels
   |
   |-----test
       |-----images
       |-----labels

确保以下几点:

所有的训练图片都位于data/train/images目录下,相应的标注文件位于data/train/labels目录下。
所有的验证图片都位于data/valid/images目录下,相应的标注文件位于data/valid/labels目录下。
所有的测试图片都位于data/test/images目录下,相应的标注文件位于data/test/labels目录下。
这样的结构使得数据的管理和模型的训练、验证和测试变得非常方便。

模型训练
 Epoch   gpu_mem       box       obj       cls    labels  img_size
 1/200     20.8G   0.01576   0.01955  0.007536        22      1280: 100%|██████████| 849/849 [14:42<00:00,  1.04s/it]
           Class     Images     Labels          P          R     [email protected] [email protected]:.95: 100%|██████████| 213/213 [01:14<00:00,  2.87it/s]
             all       3395      17314      0.994      0.957      0.0957      0.0843

 Epoch   gpu_mem       box       obj       cls    labels  img_size
 2/200     20.8G   0.01578   0.01923  0.007006        22      1280: 100%|██████████| 849/849 [14:44<00:00,  1.04s/it]
           Class     Images     Labels          P          R     [email protected] [email protected]:.95: 100%|██████████| 213/213 [01:12<00:00,  2.95it/s]
             all       3395      17314      0.996      0.956      0.0957      0.0845

 Epoch   gpu_mem       box       obj       cls    labels  img_size
 3/200     20.8G   0.01561    0.0191  0.006895        27      1280: 100%|██████████| 849/849 [10:56<00:00,  1.29it/s]
           Class     Images     Labels          P          R     [email protected] [email protected]:.95: 100%|███████   | 187/213 [00:52<00:00,  4.04it/s]
             all       3395      17314      0.996      0.957      0.0957      0.0845

5.核心代码讲解

5.2 predict.py
class DetectionPredictor(BasePredictor):
    def postprocess(self, preds, img, orig_imgs):
        preds = ops.non_max_suppression(preds,
                                        self.args.conf,
                                        self.args.iou,
                                        agnostic=self.args.agnostic_nms,
                                        max_det=self.args.max_det,
                                        classes=self.args.classes)

        if not isinstance(orig_imgs, list):
            orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)

        results = []
        for i, pred in enumerate(preds):
            orig_img = orig_imgs[i]
            pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
            img_path = self.batch[0][i]
            results.append(Results(orig_img, path=img_path, names=self.model.names, boxes=pred))
        return results

这是一个名为predict.py的程序文件,它是基于检测模型进行预测的。该文件使用了Ultralytics YOLO库。

该文件定义了一个名为DetectionPredictor的类,它继承自BasePredictor类。该类用于基于检测模型进行预测。

在该文件中,还定义了一个postprocess方法,用于对预测结果进行后处理,并返回一个Results对象的列表。

该文件还包含了一些导入的模块和函数,如BasePredictor类、Results类和ops模块。

该文件的示例用法如下:

from ultralytics.utils import ASSETS
from ultralytics.models.yolo.detect import DetectionPredictor

args = dict(model='yolov8n.pt', source=ASSETS)
predictor = DetectionPredictor(overrides=args)
predictor.predict_cli()

以上是对该程序文件的概述。

5.3 train.py
# Ultralytics YOLO , AGPL-3.0 license

from copy import copy

import numpy as np

from ultralytics.data import build_dataloader, build_yolo_dataset
from ultralytics.engine.trainer import BaseTrainer
from ultralytics.models import yolo
from ultralytics.nn.tasks import DetectionModel
from ultralytics.utils import LOGGER, RANK
from ultralytics.utils.torch_utils import de_parallel, torch_distributed_zero_first


class DetectionTrainer(BaseTrainer):
    def build_dataset(self, img_path, mode='train', batch=None):
        gs = max(int(de_parallel(self.model).stride.max() if self.model else 0), 32)
        return build_yolo_dataset(self.args, img_path, batch, self.data, mode=mode, rect=mode == 'val', stride=gs)

    def get_dataloader(self, dataset_path, batch_size=16, rank=0, mode='train'):
        assert mode in ['train', 'val']
        with torch_distributed_zero_first(rank):
            dataset = self.build_dataset(dataset_path, mode, batch_size)
        shuffle = mode == 'train'
        if getattr(dataset, 'rect', False) and shuffle:
            LOGGER.warning("WARNING ⚠️ 'rect=True' is incompatible with DataLoader shuffle, setting shuffle=False")
            shuffle = False
        workers = 0
        return build_dataloader(dataset, batch_size, workers, shuffle, rank)

    def preprocess_batch(self, batch):
        batch['img'] = batch['img'].to(self.device, non_blocking=True).float() / 255
        return batch

    def set_model_attributes(self):
        self.model.nc = self.data['nc']
        self.model.names = self.data['names']
        self.model.args = self.args

    def get_model(self, cfg=None, weights=None, verbose=True):
        model = DetectionModel(cfg, nc=self.data['nc'], verbose=verbose and RANK == -1)
        if weights:
            model.load(weights)
        return model

    def get_validator(self):
        self.loss_names = 'box_loss', 'cls_loss', 'dfl_loss'
        return yolo.detect.DetectionValidator(self.test_loader, save_dir=self.save_dir, args=copy(self.args))

    def label_loss_items(self, loss_items=None, prefix='train'):
        keys = [f'{prefix}/{x}' for x in self.loss_names]
        if loss_items is not None:
            loss_items = [round(float(x), 5) for x in loss_items]
            return dict(zip(keys, loss_items))
        else:
            return keys

    def progress_string(self):
        return ('\n' + '%11s' *
                (4 + len(self.loss_names))) % ('Epoch', 'GPU_mem', *self.loss_names, 'Instances', 'Size')

    def plot_training_samples(self, batch, ni):
        plot_images(images=batch['img'],
                    batch_idx=batch['batch_idx'],
                    cls=batch['cls'].squeeze(-1),
                    bboxes=batch['bboxes'],
                    paths=batch['im_file'],
                    fname=self.save_dir / f'train_batch{ni}.jpg',
                    on_plot=self.on_plot)

    def plot_metrics(self):
        plot_results(file=self.csv, on_plot=self.on_plot)

    def plot_training_labels(self):
        boxes = np.concatenate([lb['bboxes'] for lb in self.train_loader.dataset.labels], 0)
        cls = np.concatenate([lb['cls'] for lb in self.train_loader.dataset.labels], 0)
        plot_labels(boxes, cls.squeeze(), names=self.data['names'], save_dir=self.save_dir, on_plot=self.on_plot)


if __name__ == '__main__':
    args = dict(model='./yolov8-C2f-EMSC.yaml', data='coco8.yaml', epochs=100)
    trainer = DetectionTrainer(overrides=args)
    trainer.train()

这个程序文件是一个用于训练检测模型的程序。它使用了Ultralytics YOLO库,并继承了BaseTrainer类。以下是程序文件的主要部分:

  1. 导入所需的库和模块。
  2. 定义了一个名为DetectionTrainer的类,它继承自BaseTrainer类。
  3. 在DetectionTrainer类中定义了一些方法,包括构建数据集、构建数据加载器、预处理批次、设置模型属性等。
  4. 在主函数中,创建了一个DetectionTrainer对象,并调用其train方法开始训练。

总体来说,这个程序文件是一个用于训练检测模型的脚本,它使用了Ultralytics YOLO库提供的功能和类来构建数据集、加载数据、训练模型等。

5.4 ui.py
import cv2
import numpy as np


    def __init__(self, img_path):
        self.img = cv2.imread(img_path, cv2.IMREAD_COLOR)
        self.gray = cv2.cvtColor(self.img, cv2.COLOR_BGR2GRAY)

    def detect(self):
        # 边缘检测
        edges = cv2.Canny(self.gray, 100, 200)

        # 特征提取
        contours, _ = cv2.findContours(edges, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        areas = [cv2.contourArea(c) for c in contours]
        max_index = np.argmax(areas)
        cnt = contours[max_index]

        # 阈值分割
        mask = np.zeros_like(self.gray)
        cv2.drawContours(mask, [cnt], 0, 255, -1)
        thresh = cv2.threshold(self.gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)[1]
        mask = cv2.bitwise_and(mask, thresh)

        # 开运算闭运算
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
        opening = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
        closing = cv2.morphologyEx(opening, cv2.MORPH_CLOSE, kernel)

        # 绘制边框
        x, y, w, h = cv2.boundingRect(closing)
        cv2.rectangle(self.img, (x, y), (x + w, y + h), (0, 255, 0), 2)

        return self.img

该程序文件是一个使用PyQt5构建的图形用户界面(GUI)应用程序。它包含一个名为ui.py的文件,其中定义了一个名为MyWidget的自定义窗口类和一个名为FruitDetector的水果检测类。

MyWidget类是一个继承自QDialog的窗口类,用于显示一个提示框,其中包含一个识别结果标签、一个文本输入框和两个按钮。用户可以在文本输入框中输入内容,并点击确定按钮或重新识别按钮来触发相应的操作。

FruitDetector类是一个水果检测类,用于对输入的图像进行边缘检测、特征提取、阈值分割、开闭运算和边框绘制等操作,最终返回检测到的水果图像。

除了上述两个类,程序文件还包含了一些导入的库和函数,用于实现图像处理、模型加载和推理等功能。其中包括argparseplatformshutiltimenumpycv2torch等库。

整个程序的主要逻辑是,在MyWidget类中,用户点击确定按钮或重新识别按钮时,会触发相应的操作,包括获取文本输入框中的内容、关闭窗口、设置全局变量等。然后,在FruitDetector类中,根据输入的图像进行检测,并返回检测结果。最后,将检测结果显示在窗口中的标签和图像上。

整个程序的运行逻辑是,首先加载模型和数据集,然后打开一个图形用户界面窗口,等待用户输入图像路径。用户输入图像路径后,程序会对图像进行检测,并将检测结果显示在窗口中。用户可以点击确定按钮关闭窗口,或点击重新识别按钮重新进行水果检测。

5.5 backbone\convnextv2.py
import torch
import torch.nn as nn
import torch.nn.functional as F
from timm.models.layers import trunc_normal_, DropPath

class LayerNorm(nn.Module):
    def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(normalized_shape))
        self.bias = nn.Parameter(torch.zeros(normalized_shape))
        self.eps = eps
        self.data_format = data_format
        if self.data_format not in ["channels_last", "channels_first"]:
            raise NotImplementedError 
        self.normalized_shape = (normalized_shape, )
    
    def forward(self, x):
        if self.data_format == "channels_last":
            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
        elif self.data_format == "channels_first":
            u = x.mean(1, keepdim=True)
            s = (x - u).pow(2).mean(1, keepdim=True)
            x = (x - u) / torch.sqrt(s + self.eps)
            x = self.weight[:, None, None] * x + self.bias[:, None, None]
            return x

class GRN(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim))
        self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim))

    def forward(self, x):
        Gx = torch.norm(x, p=2, dim=(1,2), keepdim=True)
        Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
        return self.gamma * (x * Nx) + self.beta + x

class Block(nn.Module):
    def __init__(self, dim, drop_path=0.):
        super().__init__()
        self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim)
        self.norm = LayerNorm(dim, eps=1e-6)
        self.pwconv1 = nn.Linear(dim, 4 * dim)
        self.act = nn.GELU()
        self.grn = GRN(4 * dim)
        self.pwconv2 = nn.Linear(4 * dim, dim)
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()

    def forward(self, x):
        input = x
        x = self.dwconv(x)
        x = x.permute(0, 2, 3, 1)
        x = self.norm(x)
        x = self.pwconv1(x)
        x = self.act(x)
        x = self.grn(x)
        x = self.pwconv2(x)
        x = x.permute(0, 3, 1, 2)

        x = input + self.drop_path(x)
        return x

class ConvNeXtV2(nn.Module):
    def __init__(self, in_chans=3, num_classes=1000, 
                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], 
                 drop_path_rate=0., head_init_scale=1.
                 ):
        super().__init__()
        self.depths = depths
        self.downsample_layers = nn.ModuleList()
        stem = nn.Sequential(
            nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4),
            LayerNorm(dims[0], eps=1e-6, data_format="channels_first")
        )
        self.downsample_layers.append(stem)
        for i in range(3):
            downsample_layer = nn.Sequential(
                    LayerNorm(dims[i], eps=1e-6, data_format="channels_first"),
                    nn.Conv2d(dims[i], dims[i+1], kernel_size=2, stride=2),
            )
            self.downsample_layers.append(downsample_layer)

        self.stages = nn.ModuleList()
        dp_rates=[x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] 
        cur = 0
        for i in range(4):
            stage = nn.Sequential(
                *[Block(dim=dims[i], drop_path=dp_rates[cur + j]) for j in range(depths[i])]
            )
            self.stages.append(stage)
            cur += depths[i]

        self.norm = nn.LayerNorm(dims[-1], eps=1e-6)
        self.head = nn.Linear(dims[-1], num_classes)

        self.apply(self._init_weights)
        self.channel = [i.size(1) for i in self.forward(torch.randn(1, 3, 640, 640))]

    def _init_weights(self, m):
        if isinstance(m, (nn.Conv2d, nn.Linear)):
            trunc_normal_(m.weight, std=.02)
            nn.init.constant_(m.bias, 0)

    def forward(self, x):
        res = []
        for i in range(4):
            x = self.downsample_layers[i](x)
            x = self.stages[i](x)
            res.append(x)
        return res

该程序文件是一个实现了ConvNeXt V2模型的PyTorch代码。ConvNeXt V2是一个用于图像分类任务的卷积神经网络模型。该模型使用了一系列的ConvNeXt块来构建网络结构。

该程序文件中定义了以下几个类和函数:

  1. LayerNorm类:实现了支持两种数据格式(channels_last和channels_first)的LayerNorm层。

  2. GRN类:实现了全局响应归一化(GRN)层。

  3. Block类:实现了ConvNeXtV2块,包括深度卷积、LayerNorm、线性层、GELU激活函数、GRN层和线性层。

  4. ConvNeXtV2类:实现了ConvNeXt V2模型,包括多个下采样层和多个特征分辨率阶段。每个阶段由多个ConvNeXt块组成。

  5. update_weight函数:用于更新模型的权重。

  6. convnextv2_atto、convnextv2_femto、convnextv2_pico、convnextv2_nano、convnextv2_tiny、convnextv2_base、convnextv2_large和convnextv2_huge函数:分别创建不同规模的ConvNeXt V2模型,并加载预训练权重(如果提供)。

该程序文件中的代码实现了ConvNeXt V2模型的网络结构和权重加载功能。可以根据需要选择不同规模的模型,并通过加载预训练权重来进行图像分类任务。

5.6 backbone\CSwomTramsformer.py
import torch
import torch.nn as nn
import torch.nn.functional as F
from functools import partial

from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
from timm.models.helpers import load_pretrained
from timm.models.layers import DropPath, to_2tuple, trunc_normal_
from timm.models.registry import register_model
from einops.layers.torch import Rearrange
import torch.utils.checkpoint as checkpoint
import numpy as np
import time

__all__ = ['CSWin_tiny', 'CSWin_small', 'CSWin_base', 'CSWin_large']

class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x

class LePEAttention(nn.Module):
    def __init__(self, dim, resolution, idx, split_size=7, dim_out=None, num_heads=8, attn_drop=0., proj_drop=0., qk_scale=None):
        super().__init__()
        self.dim = dim
        self.dim_out = dim_out or dim
        self.resolution = resolution
        self.split_size = split_size
        self.num_heads = num_heads
        head_dim = dim // num_heads
        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
        self.scale = qk_scale or head_dim ** -0.5
        if idx == -1:
            H_sp, W_sp = self.resolution, self.resolution
        elif idx == 0:
            H_sp, W_sp = self.resolution, self.split_size
        elif idx == 1:
            W_sp, H_sp = self.resolution, self.split_size
        else:
            print ("ERROR MODE", idx)
            exit(0)
        self.H_sp = H_sp
        self.W_sp = W_sp
        stride = 1
        self.get_v = nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1,groups=dim)

        self.attn_drop = nn.Dropout(attn_drop)

    def im2cswin(self, x):
        B, N, C = x.shape
        H = W = int(np.sqrt(N))
        x = x.transpose(-2,-1).contiguous().view(B, C, H, W)
        x = img2windows(x, self.H_sp, self.W_sp)
        x = x.reshape(-1, self.H_sp* self.W_sp, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3).contiguous()
        return x

    def get_lepe(self, x, func):
        B, N, C = x.shape
        H = W = int(np.sqrt(N))
        x = x.transpose(-2,-1).contiguous().view(B, C, H, W)

        H_sp, W_sp = self.H_sp, self.W_sp
        x = x.view(B, C, H // H_sp, H_sp, W // W_sp, W_sp)
        x = x.permute(0, 2, 4, 1, 3, 5).contiguous().reshape(-1, C, H_sp, W_sp) ### B', C, H', W'

        lepe = func(x) ### B', C, H', W'
        lepe = lepe.reshape(-1, self.num_heads, C // self.num_heads, H_sp * W_sp).permute(0, 1, 3, 2).contiguous()

        x = x.reshape(-1, self.num_heads, C // self.num_heads, self.H_sp* self.W_sp).permute(0, 1, 3, 2).contiguous()
        return x, lepe

    def forward(self, qkv):
        """
        x: B L C
        """
        q,k,v = qkv[0], qkv[1], qkv[2]

        ### Img2Window
        H = W = self.resolution
        B, L, C = q.shape
        assert L == H * W, "flatten img_tokens has wrong size"
        
        q = self.im2cswin(q)
        k = self.im2cswin(k)
        v, lepe = self.get_lepe(v, self.get_v)

        q = q * self.scale
        attn = (q @ k.transpose(-2, -1))  # B head N C @ B head C N --> B head N N
        attn = nn.functional.softmax(attn, dim=-1, dtype=attn.dtype)
        attn = self.attn_drop(attn)

        x = (attn @ v) + lepe
        x = x.transpose(1, 2).reshape(-1, self.H_sp* self.W_sp, C)  # B head N N @ B head N C

        ### Window2Img
        x = windows2img(x, self.H_sp, self.W_sp, H, W).view(B, -1, C)  # B H' W' C

        return x


class CSWinBlock(nn.Module):

    def __init__(self, dim, reso, num_heads,
                 split_size=7, mlp_ratio=4., qkv_bias=False, qk_scale=None,
                 drop=0., attn_drop=0., drop_path=0.,
                 act_layer=nn.GELU, norm_layer=nn.LayerNorm,
                 last_stage=False):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.patches_resolution = reso
        self.split_size = split_size
        self.mlp_ratio = mlp_ratio
        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.norm1 = norm_layer(dim)

        if self.patches_resolution == split_size:
            last_stage = True
        if last_stage:
            self.branch_num = 1
        else:
            self.branch_num = 2
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(drop)
        
        if last_stage:
            self.attns = nn.ModuleList([
                LePEAttention(
                    dim, resolution=self.patches_resolution, idx = -1,
                    split_size=split_size, num_heads=num_heads, dim_out=dim,
                    qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
                for i in range(self.branch_num)])
        else:
            self.attns = nn.ModuleList([
                LePEAttention(
                    dim//2, resolution=self.patches_resolution, idx = i,
                    split_size=split_size, num_heads=num_heads//2, dim_out=dim//2,
                    qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
                for i in range(self.branch_num)])
        

        mlp_hidden_dim = int(dim * mlp_ratio)

        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, out_features=dim, act_layer=act_layer, drop=drop)
        self.norm2 = norm_layer(dim)

    def forward(self, x):
        """
        x: B, H*W, C
        """

        H = W = self.patches_resolution
        B, L, C = x.shape
        assert L == H * W, "flatten img_tokens has wrong size"
        img = self.norm1(x)
        qkv = self.qkv(img).reshape(B, -1, 3, C).permute(2, 0, 1, 3)
        
        if self.branch_num == 2:
            x1 = self.attns[0](qkv[:,:,:,:C//2])
            x2 = self.attns[1](qkv[:,:,:,C//2:])
            attened_x = torch.cat([x1,x2], dim=2

该程序文件是一个用于实现CSWin Transformer模型的Python文件。CSWin Transformer是一种用于图像分类任务的Transformer模型,其主要特点是使用了局部位置编码(Local Position Encoding)和窗口化输入(Window-based Input)的方法来处理图像数据。

该文件定义了以下几个类:

  1. Mlp:一个多层感知机(MLP)模块,用于对输入进行线性变换和激活函数处理。

  2. LePEAttention:一个局部位置编码(LePE)注意力模块,用于计算输入的注意力权重。

  3. CSWinBlock:CSWin Transformer的基本模块,包含了局部位置编码注意力模块和多层感知机模块。

  4. Merge_Block:用于将CSWinBlock的输出进行卷积和归一化处理。

此外,该文件还定义了一些辅助函数和常量。

总体来说,该文件实现了CSWin Transformer模型的各个组件,并提供了模型的前向传播方法。

6.系统整体结构

整体功能和构架概述:

该项目是一个视觉项目,主要用于金属工件缺陷检测系统。它使用了多个深度学习模型和算法来进行图像分类、目标检测和图像分割等任务。整个项目的构架包括了数据处理、模型训练、模型推理和图形用户界面等部分。

以下是每个文件的功能概述:

文件路径 功能概述
export.py 将YOLOv8 PyTorch模型导出为其他格式的函数
predict.py 基于检测模型进行预测的程序
train.py 用于训练检测模型的程序
ui.py 使用PyQt5构建的图形用户界面(GUI)应用程序
backbone\convnextv2.py 实现了ConvNeXt V2模型的PyTorch代码
backbone\CSwomTramsformer.py 实现了CSWin Transformer模型的Python文件
backbone\EfficientFormerV2.py 实现了EfficientFormer V2模型的PyTorch代码
backbone\efficientViT.py 实现了EfficientViT模型的PyTorch代码
backbone\fasternet.py 实现了Fasternet模型的PyTorch代码
backbone\lsknet.py 实现了LSKNet模型的PyTorch代码
backbone\repvit.py 实现了RepVIT模型的PyTorch代码
backbone\revcol.py 实现了RevCoL模型的PyTorch代码
backbone\SwinTransformer.py 实现了Swin Transformer模型的PyTorch代码
backbone\VanillaNet.py 实现了VanillaNet模型的PyTorch代码
extra_modules\afpn.py 实现了AFPN模块的PyTorch代码
extra_modules\attention.py 实现了注意力机制模块的PyTorch代码
extra_modules\block.py 实现了各种块模块的PyTorch代码
extra_modules\dynamic_snake_conv.py 实现了动态蛇形卷积模块的PyTorch代码
extra_modules\head.py 实现了模型头部模块的PyTorch代码
extra_modules\kernel_warehouse.py 实现了核函数仓库的PyTorch代码
extra_modules\orepa.py 实现了OREPA模块的PyTorch代码
extra_modules\rep_block.py 实现了RepBlock模块的PyTorch代码
extra_modules\RFAConv.py 实现了RFAConv模块的PyTorch代码
extra_modules_init_.py 额外模块的初始化文件
extra_modules\ops_dcnv3\setup.py DCNv3模块的安装文件
extra_modules\ops_dcnv3\test.py DCNv3模块的测试文件
extra_modules\ops_dcnv3\functions\dcnv3_func.py DCNv3模块的函数实现
extra_modules\ops_dcnv3\functions_init_.py DCNv3模块函数的初始化文件
extra_modules\ops_dcnv3\modules\dcnv3.py DCNv3模块的模型实现
extra_modules\ops_dcnv3\modules_init_.py DCNv3模块的初始化文件
models\common.py 包含一些通用的模型函数和类
models\experimental.py 包含一些实验性的模型函数和类
models\tf.py 包含一些TensorFlow模型函数和类
models\yolo.py 包含YOLO模型的函数和类
models_init_.py 模型的初始化文件
utils\activations.py 包含一些激活函数的实现
utils\augmentations.py 包含一些数据增强方法的实现
utils\autoanchor.py 包含自动锚框生成的实现
utils\autobatch.py 包含自动批次大小调整的实现
utils\callbacks.py 包含一些回调函数的实现
utils\datasets.py 包含数据集的处理和加载方法的实现
utils\downloads.py 包含下载文件的方法的实现
utils\general.py 包含一些通用的辅助函数的实现
utils\loss.py 包含一些损失函数的实现
utils\metrics.py 包含一些评估指标的实现
utils\plots.py 包含绘图函数的实现
utils\torch_utils.py 包含一些PyTorch工具函数的实现
utils_init_.py 工具函数的初始化文件
utils\aws\resume.py 包含AWS训练恢复的实现
utils\aws_init_.py AWS工具的初始化文件
utils\flask_rest_api\example_request.py Flask REST API示例请求的实现
utils\flask_rest_api\restapi.py Flask REST API的实现
utils\loggers_init_.py 日志记录器的初始化文件
utils\loggers\wandb\log_dataset.py 使用WandB记录数据集的实现

7.YOLOv8简介

在各种目标检测算法中,YOLO系列因其速度和准确度间的优异平衡脱颖而出,能够准确、快速的识别目标,便于部署到各种移动设备中,已经广泛应用于各种领域的目标检测、跟踪和分割。目前最新版本 YOLOv8由原v5的开发团队Ultralytics于2023年1月提出,按规格大小可划分为n、s、m、1和x5个版本,是目前最先进的目标检测算法,有着优异的性能,很适合用于无人机航拍图像目标检测。其网络结构如图所示。
【原创创新点】金属工件缺陷检测系统:Efficient Multi-Scale-Conv的改进YOLOv8_第2张图片
YOLOv8模型包括Input、Backbone、Neck 和Head4部分。其中 Input选用了Mosaic数据增强方法,并且对于不同大小的模型,有部分超参数会进行修改,典型的如大模型会开启 MixUp 和CopyPaste数据增强,能够丰富数据集,提升模型的泛化能力和鲁棒性。Backbone主要用于提取图片中的信息,提供给Neck和Head使用,由多个Conv、C2f模块和尾部的SPPF组成。Conv模块由单个Conv2d、
BatchNorm2d和激活函数构成,用丁是双行在万E5特征图;YOLOv8参考了C3模块的残左绒以心YOLOv7[16l的ELAN思想,设计出了C2f 结构,可以在保证轻量化的同时获得更加丰富的梯度流信息,并根据模型尺度来调整通道数,大幅提升了模型性能;SPPF是空间金字塔池化,能够融合不同尺度的特征。Neck部分主要起特征融合的作用,充分利用了骨干网络提取的特征,采用FPN[17]+PAN[18〕结构,
能够增强多个尺度上的语义表达和定位能力。Head输出端根据前两部分处理得到的特征来获取检测目标的类别和位置信息,做出识别,换以小旦公来和定解耦头结构,将分类和检测头分:9w‘无锅框的位关注侧重点不同的问题,同时也采用了无锚框的目标检测(Anchor-Free),能够提升检测速度。Loss计算方面采用了正负样本动态分配策略,使用 VFLLoss 作为分类损失,使用DFLLoss+CIOU Loss作为回归损失。

8.AutoFocus: Efficient Multi-Scale Conv简介

参考该博客提出了AutoFocus,一种高效的多尺度目标检测算法。相较于以前对整个图像金字塔进行处理,该方法以一种由表及里的姿态,仅处理哪些整体上看来很有可能存在小物体的区域。这个可以通过预测一张类别未知的分割图FocusPixels来得到。为了高效利用FocusPixels,另外一个算法用于产生包含FocusPixels的FocusChips,这样可以减少计算量并处理更更精细的尺度。在不同尺度上FocusChips所得到的检测结果综合时,会出现问题,我们也提供了解决问题的方案。AutoFocus在COCO上的结果有49.7%mAP(50%重叠下68.3%),与多尺度baseline相仿但是快了2.5倍。金字塔中处理的像素数量减少了5倍mAP只下降1%,在与RetinaNet采用相同的ResNet-101结构且速度相同时,高了10%mAP。

人类寻找物体是一个动态的过程,且寻找时间与场景的复杂度是直接相关的。当我们的眼神在不同的点之间漂移时,其他的区域会被我们有意忽视。然而,现在的检测算法是一种静态的推理过程且图像金字塔中的每个像素都受到了一样的对待,这使得过程变得没有效率。现在许多的应用都不强调实时性,而计算上的节省其实可以产生很大收益。

在COCO数据集中,虽然40%的物体都是小物体,但是它们一共只占了全图的0.3%。如果金字塔是3倍关系,则在高分辨率层要进行9倍的运算。XXX 。那么有没有办法在低分辨率下找出可能含有这些物体的区域呢?

以人眼方案类比,我们可以从低分辨率图像开始,找出有可能存在物体的区域再“聚集”到高分辨率。我们的AutoFocus会在一层中计算小物体的分割图FocusPixels,而在每一个FocusPixels上会用一个算法产生下一层所需要关注的chips。在COCO的最大分辨率层上我们可以只处理**20%的像素而性能不下降,如果只处理5%**也只下降1%而已。

图像金字塔与卷积神经网络对CV十分重要。然而卷积神经网络无法做到对尺度不敏感,所以为了不同大小的物体需要依赖图像金字塔。虽然训练已经有了高效的方法,但是其推断时间依然远离实际使用标准。

目标检测加速有很长历史了。常用的有特征近似以减少尺度、级联、特征金字塔,且最后一个最近很多人研究。

AutoFocus为速度与精度之间提供了一个平滑的折衷,指出可以在低分辨率下看出小物体的所在,从而节约计算。FocusPixels的计算较为容易。

先简单介绍一下SNIP,是一种多尺度的训练、推断算法。主要思想是训练针对某个特定scale的检测器而不是scale-invariant检测器。这样训练样本就局限于在某个尺度范围内,以适于这个检测器处理。比如在高分辨率仅处理小物体而高分辨率仅处理大物体,其优势在于训练时不用考虑尺度的变化。

由于训练时物体大小是已知的,我们在图像金字塔中可以忽略大量区域而只处理物体周围的区域。SNIPER说明这样低分辨率的训练与全图训练相比并不会降低性能。同样,在推断过程中如果可以在大分辨率图像上预测可能出现小物体的chip,我们也就不用处理整张高分辨率图片。在训练时,许多物体会被裁剪、扭曲,这可以当作是一种数据扩增,然而当这个情况在推断时出现,则会产生错误,所以我们还需要一个算法来整合不同尺度上的检测结果。

AutoFocus框架

如SIFT、SURF等传统分类特征包含两个组件,一个detector和一个descriptor。detector只包含轻量级的操作如DoG、LoG,用于在整张图片上找到感兴趣的区域;descriptor,一般来说计算量比较大,则只需要关注那些咸兴趣的区域。这个级联结构保障了图片处理的效率。

同样,AutoFocus框架也是用于预测感兴趣的区域,并丢弃在下一尺度不可能存在物体的区域,并将裁剪和缩放后的区域传递给下一尺度。AutoFocus由三个部分组成:FocusPixels,FocusChips和focus stacking。

FocusPixels

FocusPixels定义在卷积网络特征图的粒度之上(如conv5),如果特征图上某个像素与小物体有重叠则标注为一个FocusPixel。(小物体:面积处于一个网络的输入范围之内)。训练过程中,FocusPixels标注为正,某些与不在面积范围内的物体有重叠的像素标注为无效,其他像素标注为负。AutoFocus的训练目标是使在FocusPixels区域产生较大的激活值。
【原创创新点】金属工件缺陷检测系统:Efficient Multi-Scale-Conv的改进YOLOv8_第3张图片
如果同时多个物体与同一像素重叠,优先给正标注。我们的网络输入是512x512,然后a,b,c取值分别是5,64,90。对于太大或太小的物体,我们认为当前尺度上没有足够人信息进行判断,所以会标定为无效。整个网络结构如下图。训练时加两层带ReLU的卷积(3x3和1x1)以及一个二分类softmax来预测FocusPixels。

【原创创新点】金属工件缺陷检测系统:Efficient Multi-Scale-Conv的改进YOLOv8_第4张图片

FocusChip生成

推断过程中,我们标注输出概率大于阈值t的像素(这个参数可以控制加速比),得到一些连通域。对于每一个域,我们做一次膨胀使其包含一些有助于识别的周围信息,膨胀后相连的部分要合并。然后,我们生成一些包含连通域的chips(不同chips可能重叠,如有重叠则合并)。有些chip可能太小缺少必要信息,并使chip的大小变化很大,为了效率我们要求保证一个最小的chip size。算法流程如下:

【原创创新点】金属工件缺陷检测系统:Efficient Multi-Scale-Conv的改进YOLOv8_第5张图片

级联推断过程

我们的过程类似于图片金字塔的推断过程但是在过程中一步步去除不关注的图片区域。我们在最低分辨率上进行检测并生成chips,然后在这些chips再进行检测和生成chips。

Focus Stacking

这种级联分类器的一个问题是在chip边缘本来的大物体可能生成一些被剪裁的检测目标。在下一个尺度上,由于剪裁过了,所以会检测到小的、错误的正例False Positive。之前算法中的第二步膨胀其实就是为了保证没有物体会出现在chip的边缘。这样,当我们在局部图片chip的边缘检测到物体时,就可以把这个检测结果丢弃,即使其在SNIP的范围内。

也有一些物体在图片边缘的特殊情况。如果chip和图片共享一条边的话,我们会检查检测结果的其他边是否完全在里面,如果在,则保留它,否则丢弃。

在每一尺度得到有效检测结果后,我们可以将不同尺度上的结果进行整合,将这些结果映射到原图坐标系上。最后,非极大值抑制NMS用于聚集。网络结构如之前的图片所示。

9.系统整合

下图完整源码&数据集&环境部署视频教程&自定义UI界面
【原创创新点】金属工件缺陷检测系统:Efficient Multi-Scale-Conv的改进YOLOv8_第6张图片

参考博客《金属工件缺陷检测系统:Efficient Multi-Scale-Conv的改进YOLOv8》

10.参考文献


[1]于海涛,李健升,刘亚姣,等.基于级联神经网络的型钢表面缺陷检测算法[J].计算机应用.2023,43(1).DOI:10.11772/j.issn.1001-9081.2021111940 .

[2]孙琪翔,何宁,张敬尊,等.基于非局部高分辨率网络的轻量化人体姿态估计方法[J].计算机应用.2022,42(5).DOI:10.11772/j.issn.1001-9081.2021030512 .

[3]胡海涛,杜昊晨,王素琴,等.改进YOLOX的药品泡罩铝箔表面缺陷检测方法
[J].图学学报.2022,43(5).DOI:10.11996/JG.j.2095-302X.2022050803 .

[4]黄凤荣,李杨,郭兰申,等.基于Faster R-CNN的零件表面缺陷检测算法[J].计算机辅助设计与图形学学报.2020,(6).DOI:10.3724/SP.J.1089.2020.17981 .

[5]AI Yong-hao,XU Ke.Surface Detection of Continuous Casting Slabs Based on Curvelet Transform and Kernel Locality Preserving Projections[J].钢铁研究学报(英文版).2013,(5).80-86.

[6]佚名.Deep Metallic Surface Defect Detection: The New Benchmark and Detection Network[J].Sensors.2020,20(6).1562.DOI:10.3390/s20061562 .

[7]Ashour, Mohammed Waleed,Khalid, Fatimah,Halin, Alfian Abdul,等.Surface Defects Classification of Hot-Rolled Steel Strips Using Multi-directional Shearlet Features[J].Arabian journal for science and engineering.2019,44(4).2925-2932.DOI:10.1007/s13369-018-3329-5 .

[8]Medina Roberto,Gayubo Fernando,González-Rodrigo Luis M.,等.Automated visual classification of frequent defects in flat steel coils[J].The International Journal of Advanced Manufacturing Technology.2011,57(9-12).1087-1097.DOI:10.1007/s00170-011-3352-0 .

[9]Girshick, R.,Donahue, J.,Darrell, T.,等.Rich Feature Hierarchies for Accurate Object Detection and Semantic Segmentation[C].

[10]François Chollet.Xception: Deep Learning with Depthwise Separable Convolutions[C].

你可能感兴趣的:(YOLO,人工智能,yolov8)