Yolov5.6增加注意力機制+ByterTrack：目標檢測與跟蹤

簡介：本項目使用ESP32-CAM采集圖像上傳至上位機進行手部目標檢測與追蹤，使用了YOLOv5.6（注意力機制ECA,CBAM）+ByteTrack

博主同款迅雷鏈接：
鏈接：https://pan.xunlei.com/s/VOSO1EIzmXhBb_BIKM58cM5cA1#
提取碼：8ms6
需要付費指導的查看簡介聯系我！

項目拉取：

git clone https://github.com/ultralytics/yolov5.git

環境準備：

Python3.10：https://mirrors.aliyun.com/python-release/windows/python-3.10.8.exe

進入Yolov5根目錄創建虛擬環境：

python -m venv venv
venv\Scripts\activate #激活虛擬環境
pip install -e .      #編輯式安裝，之后修改Yolov5源碼直接運行

請添加圖片描述

數據集準備：

手部數據egohands_data：http://vision.soic.indiana.edu/egohands_files/egohands_data.zip

數據格式轉換：https://blog.csdn.net/wukong168/article/details/122783179

請添加圖片描述

需要添加自定義數據集的可以錄制視頻，使用腳本提取幀，使用Labelme標注圖像。

數據集文件夾目錄必須符合一下結構，且圖像與標注文件名必須一致：

-images-train #訓練集圖片-val   #驗證集圖片
-labels-train #訓練集標注文件-val   #驗證集圖片文件

coco128.yaml配置是數據集路徑：

path: D:/VUE/ModelData/egohands_data/data
train: images/train
val: images/val
test: # test images (optional)# Classes
names:0: Hand

YOLOv5結構：

使用AI簡單了解即可。
請添加圖片描述

注意力機制添加:

在Yolov5根目錄下的models/common.py文件末尾添加自定義模塊，也可以創建一個獨立的文件表示一個模塊：

CBAM:

請添加圖片描述

import torch
import math
import torch.nn as nn
import torch.nn.functional as Fclass BasicConv(nn.Module):def __init__(self, in_planes, out_planes, kernel_size, stride=1, padding=0, dilation=1, groups=1, relu=True, bn=True, bias=False):super(BasicConv, self).__init__()self.out_channels = out_planesself.conv = nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)self.bn = nn.BatchNorm2d(out_planes,eps=1e-5, momentum=0.01, affine=True) if bn else Noneself.relu = nn.ReLU() if relu else Nonedef forward(self, x):x = self.conv(x)if self.bn is not None:x = self.bn(x)if self.relu is not None:x = self.relu(x)return xclass Flatten(nn.Module):def forward(self, x):return x.view(x.size(0), -1)class ChannelGate(nn.Module):def __init__(self, gate_channels, reduction_ratio=16, pool_types=['avg', 'max']):super(ChannelGate, self).__init__()self.gate_channels = gate_channelsself.mlp = nn.Sequential(Flatten(),nn.Linear(gate_channels, gate_channels // reduction_ratio),nn.ReLU(),nn.Linear(gate_channels // reduction_ratio, gate_channels))self.pool_types = pool_typesdef forward(self, x):channel_att_sum = Nonefor pool_type in self.pool_types:if pool_type=='avg':avg_pool = F.avg_pool2d( x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))channel_att_raw = self.mlp( avg_pool )elif pool_type=='max':max_pool = F.max_pool2d( x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))channel_att_raw = self.mlp( max_pool )elif pool_type=='lp':lp_pool = F.lp_pool2d( x, 2, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))channel_att_raw = self.mlp( lp_pool )elif pool_type=='lse':# LSE pool onlylse_pool = logsumexp_2d(x)channel_att_raw = self.mlp( lse_pool )if channel_att_sum is None:channel_att_sum = channel_att_rawelse:channel_att_sum = channel_att_sum + channel_att_rawscale = F.sigmoid( channel_att_sum ).unsqueeze(2).unsqueeze(3).expand_as(x)return x * scaledef logsumexp_2d(tensor):tensor_flatten = tensor.view(tensor.size(0), tensor.size(1), -1)s, _ = torch.max(tensor_flatten, dim=2, keepdim=True)outputs = s + (tensor_flatten - s).exp().sum(dim=2, keepdim=True).log()return outputsclass ChannelPool(nn.Module):def forward(self, x):return torch.cat( (torch.max(x,1)[0].unsqueeze(1), torch.mean(x,1).unsqueeze(1)), dim=1 )class SpatialGate(nn.Module):def __init__(self):super(SpatialGate, self).__init__()kernel_size = 7self.compress = ChannelPool()self.spatial = BasicConv(2, 1, kernel_size, stride=1, padding=(kernel_size-1) // 2, relu=False)def forward(self, x):x_compress = self.compress(x)x_out = self.spatial(x_compress)scale = F.sigmoid(x_out) # broadcastingreturn x * scaleclass CBAM(nn.Module):def __init__(self, gate_channels, reduction_ratio=16, pool_types=['avg', 'max'], no_spatial=False):super(CBAM, self).__init__()self.ChannelGate = ChannelGate(gate_channels, reduction_ratio, pool_types)self.no_spatial=no_spatialif not no_spatial:self.SpatialGate = SpatialGate()def forward(self, x):x_out = self.ChannelGate(x)if not self.no_spatial:x_out = self.SpatialGate(x_out)return x_out

請添加圖片描述

class ECA(nn.Module):"""Constructs a ECA module.Args:channel: Number of channels of the input feature mapk_size: Adaptive selection of kernel size"""def __init__(self, c1, c2, k_size=3):super(ECA, self).__init__()self.avg_pool = nn.AdaptiveAvgPool2d(1)self.conv = nn.Conv1d(1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False)self.sigmoid = nn.Sigmoid()def forward(self, x):# feature descriptor on the global spatial informationy = self.avg_pool(x)y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)# Multi-scale information fusiony = self.sigmoid(y)return x * y.expand_as(x)

添加到YOLO系統配置輸入參數：

在Yolov5根目錄下的models/yolo.py文件導入自定義模塊：
請添加圖片描述
輸入參數配置，此處有大量if語句用于個性化輸入通道配置：

修改YOLOv5結構文件：

在Yolov5根目錄下的models/yolov5s.yaml文件定義模型結構：

#YOLOv5s_custom summary: 320 layers, 1046153 parameters, 1046153 gradients, 2.5 GFLOPs
#YOLOv5s summary:        214 layers, 7022326 parameters, 7022326 gradients, 15.9 GFLOPs
nc: 1 # number of classes
depth_multiple: 0.33 # model depth multiple
width_multiple: 0.25 # layer channel multiple
anchors:- [10, 13, 16, 30, 33, 23] # P3/8- [30, 61, 62, 45, 59, 119] # P4/16- [116, 90, 156, 198, 373, 326] # P5/32activation: nn.RReLU() # 激活函數可以自定義# YOLOv5 v6.0 backbone
backbone:# [from, number, module, args][[-1, 1, Conv, [64, 3, 2, 1]], # 0-P1/2[-1, 1, Conv, [128, 3, 2]], # 1-P2/4[-1, 3, C3, [128]], # 2[-1, 1, Conv, [256, 3, 2]], # 3-P3/8[-1, 6, C3, [256]],[-1, 1, Conv, [512, 3, 2]], # 5-P4/16[-1, 4, C3, [512]],[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32[-1, 1, C3, [1024]],[-1, 1, ECA, [1024]],  # 添加 ECA 提高小目標檢測效果[-1, 1, SPPF, [1024, 5]], # 10]# YOLOv5 v6.0 head
head: [# [from, number, module, args][-1, 1, GhostConv, [512, 1, 1]], # 11[-1, 1, nn.Upsample, [None, 2, "nearest"]],[[-1, 6], 1, Concat, [1]], #13 cat backbone P4[-1, 3, C3, [512, False]], # 14[-1, 1, GhostConv, [256, 1, 1]],[-1, 1, nn.Upsample, [None, 2, "nearest"]],[[-1, 4], 1, Concat, [1]], # 17 cat backbone P3[-1, 3, C3, [256, False]], # 18 (P3/8-small)[-1, 1, ECA, [256]],  # 在檢測頭中添加 ECA[-1, 1, GhostConv, [256, 3, 2]],[[-1, 14], 1, Concat, [1]], # 21 cat head P4[-1, 3, C3, [512, False]], # 22 (P4/16-medium)[-1, 1, GhostConv, [512, 3, 2]],[[-1, 10], 1, Concat, [1]], # cat head P5[-1, 3, C3, [128, False]], # 25 (P5/32-large)[[18, 22, 25], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)]

訓練：

? 由于修改了模型結構，建議不加載官方預訓練權重yolov5s.pt從0訓練，當然可以提前使用遺傳算法獲取合適的超參數加快收斂速度。

python train.py --data coco128.yaml --weights '' --cfg yolov5s.yaml --img 640

ByteTrack引入：

原理：Yolov5模型輸出檢測框列表，ByteTrack接收預測跟蹤目標，最重要的是ID參數。

具體查看博主重新編寫的檢測函數：

# -*- coding: utf-8 -*-
"""
@FileName：YoloModelRun.py
@Description：
@Author： 核知坊,一個激發創造力的網站:http://wwww.CoreKSets.cn 
@Time：2025/3/23 20:02
"""
import pathlib
import time
import numpy as np
from PIL import Image
import torch
from models.common import DetectMultiBackend
from utils.augmentations import letterbox
from utils.general import non_max_suppression
from ultralytics.utils.plotting import Annotator
from utils.dataloaders import LoadImages
from utils.general import (LOGGER,Profile,check_img_size,scale_boxes,)
from bytetracker import BYTETracker# 兼容 Windows 下的路徑問題
temp = pathlib.PosixPath
pathlib.PosixPath = pathlib.WindowsPathclass YoloModelRun():"""@ClassName：YoloModelRun@Description：@Author：錦沐Python"""def __init__(self, weights="best.onnx", imgsz = (640, 640), conf_thres=0.2, iou_thres=0.45, save_result = False, half=False):self.weights = weightsself.conf_thres = conf_thresself.iou_thres = iou_thresself.save_result = save_resultself.device = torch.device("cpu")# Load modelLOGGER.info("模型加載中")self.Model = DetectMultiBackend(weights, device=self.device, dnn=False, fp16=half)self.Model.eval()# 追蹤系統self.Tracker = BYTETracker(track_thresh=0.25, track_buffer=10, match_thresh=0.7, frame_rate=8)self.stride=self.Model.strideself.names=self.Model.namesself.pt = self.Model.ptself.imgsz = check_img_size(imgsz, s=self.stride)  # check image sizeself.Model.warmup(imgsz=(1 , 3, *self.imgsz))  # warmup# [[x1, y1, x2, y2, conf, cls], [x1, y1, x2, y2, conf, cls],....]self.detect_results = []LOGGER.info("模型準備就緒")def detect_image(self,source)->(list[[float,float,float,float,float,int,int]],Image.Image):self.detect_results.clear()# 處理不同類型的輸入if isinstance(source, str):# 如果輸入是圖片路徑，使用 LoadImages 進行加載dataset = LoadImages(source, img_size=self.imgsz, stride=self.stride, auto=self.pt, vid_stride=1)path, im, im0s, vid_cap, s = next(iter(dataset))  # 讀取一張圖片elif isinstance(source, Image.Image):im0s = np.array(source)  # 原始圖像 (RGB)im0s = im0s[..., ::-1]  # 轉換為 BGR 格式（因為 OpenCV 使用 BGR）# **使用與 LoadImages 類似的預處理**im = letterbox(im0s, self.imgsz, stride=self.stride, auto=self.pt)[0]  # 調整大小im = im.transpose((2, 0, 1))[::-1]  # HWC -> CHW, BGR -> RGBim = np.ascontiguousarray(im)  # 確保數據連續else:raise ValueError("輸入必須是圖片路徑 (str) 或 PIL.Image 對象！")# Run inferencedetect_time_tuble = (Profile(device=self.device), Profile(device=self.device), Profile(device=self.device), Profile(device=self.device))# 圖像預處理with detect_time_tuble[0]:im = torch.from_numpy(im).to(self.Model.device)im = im.half() if self.Model.fp16 else im.float()  # uint8 to fp16/32im /= 255  # 0 - 255 to 0.0 - 1.0if len(im.shape) == 3:im = im[None]  # expand for batch dim# Inference 預測with detect_time_tuble[1]:pred = self.Model(im, augment=False, visualize=False)# NMSwith detect_time_tuble[2]:pred = non_max_suppression(pred, self.conf_thres, self.iou_thres, None, False, max_det=10)# 目標繪制im0 = im0s.copy()annotator = Annotator(im0, line_width=3, example=str(self.names[0]))# Process predictionswith detect_time_tuble[3]:det = pred[0]# print(det)s = ""if len(det):det = self.Tracker.update(det)# 沒有跟蹤到物體直接返回if not det:return  [], Nonedet = torch.tensor(det, dtype=torch.float32)# print(det)# Rescale boxes from img_size to im0 sizedet[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()# Print resultsfor c in det[:, 5].unique():n = (det[:, 5] == c).sum()  # detections per classs += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, "  # add to stringfor *xyxy, conf, cls, id in reversed(det):x1,y1,x2,y2 = xyxyc = int(cls)  # integer classlabel = f"{self.names[c]} {conf:.2f}"annotator.box_label(xyxy, label, color= (255, 36, 125))# 添加檢測結果w = 640h = 480self.detect_results.append([round(x1.item() / w, 6),  # 歸一化 x1，保留6位小數round(y1.item() / h, 6),  # 歸一化 y1，保留6位小數round(x2.item() / w, 6),  # 歸一化 x2，保留6位小數round(y2.item() / h, 6),  # 歸一化 y2，保留6位小數round(conf.item(), 6),  # 置信度，保留6位小數self.names[c],  # 類別名稱int(id)])# Stream resultsif self.save_result:timestamp = int(time.time())annotator.save(f"{timestamp}_result.jpg")show_img = Image.fromarray(np.asarray(annotator.im)[..., ::-1])# Print time (inference-only)LOGGER.info(f"{s}{'' if len(det) else '(沒有檢測到目標), '}{detect_time_tuble[1].dt * 1e3:.1f}ms")# Print resultsLOGGER.info(f"Speed: %.1fms 預處理, %.1fms 識別, %.1fms NMS ,%d.1ms ByteTrack 每張圖像 {(1, 3, *self.imgsz)}" %(detect_time_tuble[0].dt * 1e3,detect_time_tuble[1].dt * 1e3,detect_time_tuble[2].dt * 1e3,detect_time_tuble[3].dt * 1e3))return self.detect_results, show_img