Yolo5,pytorch数字仪表读数,准确率高达99%,自己所写算法。
最近用YOLOv5做了个超有意思的项目——工业数字仪表自动读数。这玩意儿在电力巡检场景实测准确率干到了99.2%,比之前用的传统方案强了不是一星半点。今天就跟大伙唠唠这个项目的实现细节,特别是自己魔改的那些关键代码。
先说数据这事儿,工业仪表最大的坑就是反光。我在数据增强里加了条特制的反射增强:
class SpecularReflection(albumentations.ImageOnlyTransform): def apply(self, img, **params): if random.random() < 0.3: h,w = img.shape[:2] kernel_size = random.randint(30,80) mask = np.zeros((h,w)) cv2.ellipse(mask, (w//2,h//2), (w//4,h//4), 0,0,360,255,-1) glow = cv2.GaussianBlur(mask, (kernel_size,kernel_size), 0) img = cv2.addWeighted(img, 1.0, cv2.cvtColor(glow,cv2.COLOR_GRAY2BGR), 0.6, 0) return img这个变换在30%的图片里随机生成椭圆光斑,模拟金属表盘的反光效果。实测能让模型在强反光场景下的识别率提升18%左右。
模型结构上把YOLOv5s的neck改成了带SE注意力的双路径结构:
class DualPath(nn.Module): def __init__(self, c1, c2): super().__init__() self.cv1 = Conv(c1, c2//2, 3) self.cv2 = Conv(c1, c2//2, 3, dilation=2) self.attn = SEBlock(c2) def forward(self, x): return self.attn(torch.cat([self.cv1(x), self.cv2(x)], 1))左边是常规卷积抓局部特征,右边用空洞卷积捕捉数字间的间距特征,最后接SE通道注意力加权。改完在密集数字场景的误检率从3.1%降到0.7%。
Yolo5,pytorch数字仪表读数,准确率高达99%,自己所写算法。
后处理部分魔改了非极大抑制:
def dynamic_nms(detections): sorted_ids = torch.argsort(detections[:,4], descending=True) keep = [] while sorted_ids.size(0): keep_id = sorted_ids[0] keep.append(keep_id) other_ids = sorted_ids[1:] overlap = bbox_iou(detections[keep_id][:4], detections[other_ids][:4]) iou_thres = 0.6 - 0.2 * (detections[keep_id][4] - 0.8) / 0.2 suppressed = torch.where(overlap > iou_thres)[0] sorted_ids = other_ids[~suppressed] return torch.tensor(keep)这个动态阈值策略让低置信度的检测框用更宽松的IOU阈值,有效解决了数字粘连问题。调试时发现0.2这个系数最合适,改大改小都会影响召回率。
读数计算用了个挺有意思的加权算法:
def calc_value(digits): positions = [d[0] for d in digits] confidences = [d[1] for d in digits] sorted_digits = [d[2] for d in sorted(zip(positions, digits), key=lambda x:x[0])] total = 0 weight_sum = 0 for i, (digit, conf) in enumerate(sorted_digits): weight = conf * (10 ** (len(sorted_digits)-i-1)) # 位数权重 total += digit * weight weight_sum += weight return total / (weight_sum / 10**(len(sorted_digits)-1)) # 归一化每个数字的置信度不仅影响自身可信度,还参与计算其所在数位的权重。比如高置信度的百位数会比低置信度的十位数更有话语权,这在部分数字被遮挡时特别管用。
踩过最大的坑是标注格式——小数点必须单独标注!刚开始把"12.3"标成[1,2,.,3]四个框,结果模型总把小数点和数字连起来认。后来改成将"12.3"整体标注并在属性里标记小数点位置才解决。
这个项目让我深刻体会到:工业场景的99%准确率比学术数据集的99.9%难得多。光是处理表盘污渍、水渍、反光这些幺蛾子就耗了两个月。不过看到机械臂终于能准确读取压力表数据时,真觉得值了。