基于YOLO的人群计数实现：从密度图回归到推理后处理全链路改造指南

简介这份资源面向深度学习与计算机视觉方向的学习者和开发者提供一套基于YOLO模型实现人群计数的完整工程方案可用于车站、商场、体育场等密集场景下的实时人数统计与资源调度。压缩包共35个文件约50KB以18个Python脚本为核心配合9个Markdown说明文档、3个YAML配置、1个Jupyter Notebook及若干辅助文件覆盖模型训练、损失函数、指标评估、数据转换与推理部署等环节。资源围绕ShanghaiTech数据集展开包含训练计划、本地与云端训练指南、环境配置脚本以及伪标注初始化、课程排序、在线真值更新等模块并针对重叠边界框去重与计数精度优化给出实现思路。目前已有79人学习下载适合希望从零搭建人群计数流水线、理解YOLO回归检测原理并掌握训练调优与轻量化部署的读者参考。1. 人群计数为什么总在“数不准”和“跑不动”之间反复横跳基于YOLO的人群计数实现本质上不是把YOLO当检测器去框出每一个人而是把“人头”当成一类密集小目标用YOLO的回归头直接预测密度图或点坐标。这个思路在近两年的密集场景方案里越来越常见原因很直接传统检测框在人群重叠超过60%时NMS会互相抑制而密度图回归天然容忍遮挡。但真正落地时大部分人卡在两个极端——要么模型在公开测试集上MAE看着还行一到自己场景就崩要么推理帧率从30掉到3连实时预览都撑不住。这篇笔记面向的是已经跑通过YOLO基础训练、想把它迁移到人群计数任务的工程师我会把数据标注格式、损失函数改造、密度图生成、推理后处理这条链路拆开讲每一步都给出可复现的命令和参数。如果你手头正好有一个“基于YOLO的人群计数实现”的代码包但不确定里面哪些模块能直接复用、哪些必须重写下面的内容可以当作一份对照检查清单。2. 把YOLO改造成密度图回归器从标注到损失函数的四个改造点2.1 标注格式的选择点标注、框标注还是密度图人群计数任务里标注格式直接决定后面损失函数怎么写。常见做法有三种一是纯点标注每个人头一个坐标点适合极度密集场景二是框标注沿用YOLO的txt格式适合人头尺寸差异大、需要区分远近的场景三是直接生成密度图用高斯核在点标注上卷积得到连续密度分布。我一般会推荐点标注加自适应高斯核的方案原因是框标注在密集区域标注成本极高而且框与框之间的IoU计算在重叠严重时几乎失去意义。如果你拿到的代码包里已经有人群标注文件先检查它的格式。假设是点标注通常每行是x y两个归一化坐标没有宽高。这时候YOLO原本的锚框机制就完全用不上了需要把检测头换成密度图回归头。具体来说把YOLO的最后一个卷积层输出通道从(5类别数)*锚框数改成1表示每个像素位置的密度值。下面是一个替换检测头的代码片段import torch import torch.nn as nn class DensityHead(nn.Module): def __init__(self, in_channels512): super().__init__() # 两层卷积把特征图压缩成单通道密度图 self.conv1 nn.Conv2d(in_channels, 256, kernel_size3, padding1) self.relu nn.ReLU(inplaceTrue) self.conv2 nn.Conv2d(256, 1, kernel_size1) # 用Softplus保证输出非负密度值不能为负 self.softplus nn.Softplus() def forward(self, x): x self.relu(self.conv1(x)) x self.conv2(x) return self.softplus(x)这段代码的关键在于最后用Softplus而不是ReLU。ReLU在零点不可导且负区梯度为零密度图回归时容易导致大片区域输出死零。Softplus在负区仍有微小梯度训练初期更稳定。参数上in_channels要和YOLO主干输出的通道数对齐常见YOLOv5s是512YOLOv8n是256改之前先打印一下特征图维度。2.2 密度图真值生成高斯核半径怎么定点标注转密度图是人群计数里最容易被忽视的一步。高斯核半径选大了密度图过度平滑模型学不到局部峰值选小了密度图稀疏损失函数震荡。常见做法是用几何自适应核根据每个人头到最近邻的距离动态调整半径。下面是一个生成密度图真值的函数import numpy as np from scipy.ndimage import gaussian_filter def generate_density_map(points, img_shape, k3): points: list of (x, y) 归一化坐标 img_shape: (h, w) 输出密度图尺寸 k: 近邻数量用于估计自适应半径 h, w img_shape density np.zeros((h, w), dtypenp.float32) if len(points) 0: return density pts np.array(points) * np.array([w, h]) # 计算每个点到最近k个点的平均距离 from scipy.spatial import cKDTree tree cKDTree(pts) dists, _ tree.query(pts, kmin(k1, len(pts))) # 去掉自身距离取平均 avg_dist dists[:, 1:].mean(axis1) # 半径设为平均距离的0.3倍经验值 sigma np.clip(avg_dist * 0.3, 1.0, 10.0) for (x, y), s in zip(pts, sigma): # 在点周围3倍sigma范围内加高斯 x0, y0 int(max(0, x-3*s)), int(max(0, y-3*s)) x1, y1 int(min(w, x3*s1)), int(min(h, y3*s1)) if x1 x0 or y1 y0: continue xx, yy np.meshgrid(np.arange(x0, x1), np.arange(y0, y1)) gauss np.exp(-((xx-x)**2 (yy-y)**2) / (2*s*s)) density[y0:y1, x0:x1] gauss # 归一化使总人数等于点数 if density.sum() 0: density density / density.sum() * len(points) return density这里k3是近邻数量太小半径估计不稳太大在稀疏区域会过度平滑。sigma的系数0.3是我在几个密集场景里试出来的如果你们场景人头特别小可以降到0.2如果人头大且稀疏可以升到0.4。注意最后一步归一化保证密度图积分等于人数否则损失函数尺度会随人数变化训练不稳定。2.3 损失函数MSE、SSIM还是组合损失密度图回归最常用的损失是MSE但它有个问题对密度高的区域和密度低的区域一视同仁导致模型在稀疏背景上浪费容量。改进方案是加SSIM结构相似性损失让模型关注局部密度分布的结构。我一般用0.7倍MSE加0.3倍SSIM的组合下面是对应实现import torch import torch.nn.functional as F def ssim_loss(pred, target, window_size11): # 简化版SSIM用平均池化代替高斯窗口 mu_pred F.avg_pool2d(pred, window_size, stride1, paddingwindow_size//2) mu_target F.avg_pool2d(target, window_size, stride1, paddingwindow_size//2) sigma_pred F.avg_pool2d(pred**2, window_size, 1, window_size//2) - mu_pred**2 sigma_target F.avg_pool2d(target**2, window_size, 1, window_size//2) - mu_target**2 sigma_cross F.avg_pool2d(pred*target, window_size, 1, window_size//2) - mu_pred*mu_target C1, C2 0.01**2, 0.03**2 ssim ((2*mu_pred*mu_target C1)*(2*sigma_cross C2)) / \ ((mu_pred**2 mu_target**2 C1)*(sigma_pred sigma_target C2)) return 1 - ssim.mean() def combined_loss(pred, target): mse F.mse_loss(pred, target) ssim ssim_loss(pred, target) return 0.7 * mse 0.3 * ssimwindow_size11是SSIM的经典值对应约11像素的局部区域。如果你们密度图分辨率低于128可以降到7。组合权重0.7/0.3不是绝对的如果发现预测密度图过于模糊把SSIM权重提到0.4如果训练早期loss下降太慢先把SSIM降到0.1等MSE稳定后再加回来。2.4 训练配置学习率、批大小和冻结策略改造完模型和损失训练配置也要跟着调。YOLO原本的检测任务学习率通常设0.01但密度图回归对学习率更敏感我一般从0.001开始用余弦退火降到0.0001。批大小受显存限制但密度图回归的批内样本方差比检测小批大小可以比检测任务小一半。冻结策略上主干前几层可以冻结但Neck部分一定要解冻因为密度图需要多尺度特征融合。# 假设用YOLOv5框架修改train.py参数 python train.py \ --data crowd.yaml \ --cfg models/yolov5s_density.yaml \ --weights yolov5s.pt \ --batch-size 8 \ --epochs 100 \ --lr0 0.001 \ --lrf 0.1 \ --freeze 5 \ --hyp data/hyp_crowd.yaml--freeze 5表示冻结主干前5层--lrf 0.1是最终学习率是初始的0.1倍。hyp_crowd.yaml里要把损失权重改成上面组合损失的比例同时把数据增强里的马赛克增强关掉因为马赛克拼接会破坏密度图的连续性。3. 推理与后处理从密度图到人数和坐标的落地链路3.1 密度图积分与峰值检测模型输出密度图后人数就是密度图所有像素值之和。但如果你还需要每个人头的坐标就要在密度图上做峰值检测。常见做法是用最大池化找局部极大值再设阈值过滤。下面是一个后处理函数import torch import torch.nn.functional as F def density_to_points(density, threshold0.1, kernel_size5): density: (1, 1, H, W) 密度图 返回: list of (x, y, score) # 最大池化找局部峰值 pooled F.max_pool2d(density, kernel_size, stride1, paddingkernel_size//2) peaks (density pooled) (density threshold) # 获取峰值坐标 coords torch.nonzero(peaks[0, 0], as_tupleFalse) scores density[0, 0][peaks[0, 0]] points [(int(x), int(y), float(s)) for (y, x), s in zip(coords, scores)] return pointsthreshold0.1是密度值阈值低于它的峰值视为噪声。kernel_size5是局部窗口大小如果人头特别密集可以降到3但会引入更多重复检测。注意这里没有做NMS因为密度图峰值本身已经经过平滑重复峰值的概率比检测框低很多。3.2 多尺度推理与切片策略高分辨率图像直接缩放到模型输入尺寸会丢失小目标。常见做法是切片推理把原图切成有重叠的子图分别推理后把密度图拼回去。切片大小一般取模型输入尺寸的1.5倍重叠区域取切片尺寸的20%。下面是一个切片推理的伪代码def slice_inference(model, image, slice_size640, overlap128): h, w image.shape[:2] density_map np.zeros((h, w), dtypenp.float32) count_map np.zeros((h, w), dtypenp.float32) stride slice_size - overlap for y in range(0, h, stride): for x in range(0, w, stride): y1, x1 min(yslice_size, h), min(xslice_size, w) y0, x0 max(0, y1-slice_size), max(0, x1-slice_size) patch image[y0:y1, x0:x1] # 预处理和推理 density model(patch) # 拼回原图重叠区域取平均 density_map[y0:y1, x0:x1] density count_map[y0:y1, x0:x1] 1 density_map density_map / np.maximum(count_map, 1) return density_mapslice_size640对应常见YOLO输入尺寸overlap128是重叠像素数。重叠区域取平均而不是取最大是为了避免切片边界处的密度突变。如果你们场景人头特别小可以把slice_size提到1280但显存占用会翻倍。3.3 人数校准线性回归修正系统偏差模型在测试集上往往有系统性偏差比如整体高估或低估。一个简单有效的校准方法是在验证集上拟合一个线性回归真实人数 a * 预测人数 b。下面是一段校准代码import numpy as np from sklearn.linear_model import LinearRegression # 在验证集上收集预测人数和真实人数 pred_counts [] # 模型预测的人数 true_counts [] # 标注的真实人数 # 拟合线性回归 reg LinearRegression().fit(np.array(pred_counts).reshape(-1, 1), true_counts) a, b reg.coef_[0], reg.intercept_ print(f校准参数: a{a:.4f}, b{b:.4f}) # 推理时应用校准 def calibrate(pred_count, a, b): return max(0, a * pred_count b)这个校准对MAE的改善通常在5%到15%之间尤其是当训练集和测试集场景差异较大时。注意校准参数要在验证集上拟合不能直接用测试集否则就是数据泄露。4. 避坑与排查人群计数落地时最容易翻车的五个地方4.1 密度图全零或全一Softplus饱和与归一化错误现象是训练几个epoch后模型输出密度图要么全是接近零的值要么全是接近一的值人数预测完全失效。原因通常是两个一是Softplus在输入很大时输出接近线性梯度消失二是密度图真值归一化时除以了零。解决方法是检查真值生成函数里density.sum()是否为零加一个极小值保护同时把Softplus换成F.relu加一个小的epsilon或者用torch.log1p(torch.exp(x))手动实现并限制输入范围。4.2 人数预测随图像分辨率线性变化现象是同一张图缩放到不同尺寸预测人数跟着缩放比例变。原因是密度图没有做分辨率归一化模型学到的密度值和像素面积相关。解决方法是在生成真值时密度图积分始终归一化到人数推理时如果输入尺寸和训练尺寸不同要对输出密度图乘以(训练尺寸/推理尺寸)^2的缩放因子。4.3 密集区域峰值合并导致漏数现象是人群最密集的区域峰值检测只找到一个点但实际有十几个人。原因是高斯核半径过大密度图峰值被平滑成一个宽峰。解决方法是在后处理时对密度图做反卷积或使用更小的峰值检测窗口同时在训练时降低高斯核半径系数从0.3降到0.2。4.4 训练loss震荡不收敛现象是loss在0.5到5之间反复跳没有下降趋势。原因可能是批大小太小导致梯度噪声大或者学习率太高。解决方法是把批大小翻倍同时学习率减半如果显存不够用梯度累积模拟大批大小。另外检查数据加载器是否打乱了顺序人群计数对样本顺序敏感连续多张密集图会导致梯度爆炸。4.5 推理速度不达标现象是模型在GPU上单张推理超过200ms达不到实时。原因是密度图回归头通道数太多或者切片推理重叠太大。解决方法是把密度头第一层卷积通道从256降到128切片重叠从20%降到10%同时用TensorRT或ONNX Runtime做图优化。实测YOLOv5s改密度头后在主流GPU上单张640x640推理可以压到15ms以内。5. 进阶技巧用不确定性估计判断哪些区域值得信密度图回归有一个天然优势模型可以在每个像素输出一个方差表示对该位置密度估计的不确定性。实现方式是把密度头输出通道改成2一个通道预测均值一个预测方差损失函数用高斯负对数似然。下面是对应实现class DensityHeadWithUncertainty(nn.Module): def __init__(self, in_channels512): super().__init__() self.conv1 nn.Conv2d(in_channels, 256, 3, padding1) self.relu nn.ReLU(inplaceTrue) # 输出两个通道均值和对数方差 self.conv2 nn.Conv2d(256, 2, 1) def forward(self, x): x self.relu(self.conv1(x)) x self.conv2(x) mu, log_var x[:, 0:1], x[:, 1:2] # 均值用Softplus保证非负方差用exp保证正 return F.softplus(mu), torch.exp(log_var) def gaussian_nll_loss(mu, log_var, target): # 高斯负对数似然 precision torch.exp(-log_var) loss 0.5 * (precision * (target - mu)**2 log_var) return loss.mean()训练时用gaussian_nll_loss替代MSE推理时除了输出人数还能输出每个像素的置信区间。实际用的时候我会把方差高的区域标出来如果一张图里高方差区域占比超过30%就提示“该区域计数结果仅供参考”。这个技巧在安防和客流统计场景里很有用因为运营方需要知道哪些数据可以直接用于决策哪些需要人工复核。校准和不确定性估计可以叠加使用先用线性回归修正系统偏差再用方差过滤掉不可信区域。我自己的习惯是每次上线新场景前先跑一遍验证集把校准参数和方差阈值记在一个配置文件里推理时直接加载。这样换场景只需要重新拟合两个参数不用重新训练模型。希望帮到你。本文还有配套的精品资源点击获取