【工业机器视觉】基于深度学习的水表盘读数识别(3-数据标注与转换)-CSDN博客
读数修正
上一篇我们已经正确对图片进行了预测,并且可视化了预测结果,接下来我们要对预测结果数据进行解析、处理、修正,得到最终读数。这里不再介绍上位机开发,本文是基于Pyside6开发的上位机。
模型加载
分别加载3个训练好的模型:
python
self.modelPointer = YOLO(model_path + weight_pointer)
self.modelSeg = YOLO(model_path + weight_seg)
self.modelDigit = YOLO(model_path + weight_digit)
modelPointer:指针目标检测
modelSeg:梅花针分割
modelDigit:字轮数字目标检测
指针预测
python
results = self.modelPointer.predict(frame, imgsz=640, verbose=False)
w_pointers = []
low_pointer = None
for result in results:
names = result.names
boxes = result.boxes
box_datas = boxes.data.cpu().numpy()
if len(box_datas) == 0:
self.warnHint('指针区域:未发现检测目标')
return
for box in box_datas:
name = names[int(box[5])] # 预测标签
if name.find('p') == 0:
# 一个低位指针 当指针个数为3个时有效
if self.pointerCount == 3:
low_pointer = box
elif name.find('area_p') == 0:
# 指针区域
w_pointers.append(box)
if len(w_pointers) == 0:
self.warnHint('未检测到有效的指针区域')
return
得到最低位指针目标和其他指针区域目标数据,然后处理重复预测的BOX,一个量程只保留一个预测目标:
python
def restore_boxes(dst_areas):
'''
以X轴从大到小排序,并且两两之间如果小于10,则取得分最大的一个box
:param dst_areas: 所有预测出来的目标区域
:return:
'''
dst_areas_s = sorted(dst_areas, key=lambda box: box[0], reverse=True)
tmp_list = []
for i in range(len(dst_areas_s) - 1):
diff = abs(dst_areas_s[i][0] - dst_areas_s[i + 1][0])
if diff <= 10:
if dst_areas_s[i][4] < dst_areas_s[i + 1][4]:
tmp_list.append(i)
else:
tmp_list.append(i + 1)
dst_boxes = []
for i, dst_box in enumerate(dst_areas_s):
if i not in tmp_list:
dst_boxes.append(dst_box)
return dst_boxes
最后返回的数据就是最终我们需要处理的目标数据。下面开始逐位处理:
python
p_flows = [0, 0, 0] # 初始化
# 截取预测出来的指针区域矩形框
top_boxes = [[top_box[1] - 5, top_box[0] - 5, top_box[3] + 5, top_box[2] + 5] for top_box in pointer_boxes]
# 3个指针区域
low_flow = int(names[int(low_pointer[5])].replace('p', '')) # 这是x0.001位的值,直接预测出来的
p_flows[2] = low_flow
# x0.01角度计算、修正
angle_flow1, msg = get_real_flow(top_boxes[1], frame, self.modelSeg, p_flows[2], 'x0.01', debug_show)
if angle_flow1 is None:
self.warnHint(msg)
return
p_flows[1] = angle_flow1
# x0.1角度计算、修正
angle_flow2, msg = get_real_flow(top_boxes[0], frame, self.modelSeg, p_flows[1], 'x0.1', debug_show)
if angle_flow2 is None:
self.warnHint(msg)
return
p_flows[0] = angle_flow2
result_pointer = get_flow(p_flows)
从低位到高位,依次进行角度计算、修正后,得到最终结果。
指针读数
定义一个名为 get_real_flow
的函数,该函数用于获取修正后的流量信息。它接受多个参数,并返回一个修正后的流量值和一个可能的错误消息。
python
def get_real_flow(box, frame, model_seg, low_flow, position, debug_show):
box
: 表示一个边界框,通常用来定位图像中的特定区域。frame
: 代表一帧视频或图像。model_seg
: 可能是一个用于分割或者识别的模型。low_flow
: 可能是低流量的一个预估值或者是阈值。position
: 流量计的位置标识。debug_show
: 一个标志位,用来指示是否显示调试信息。
接下来,设置了一些固定的字符串值:
python
fix_up_value = '0,1,2,3'
fix_down_value = '7,8,9'
用于校正流量读数的固定值,与刻度盘上的某些数值相关联。然后,计算了矩形框的位置,并从原始帧中裁剪出这个矩形区域:
python
top, left, bottom, right = get_rect(box, frame)
frame_rect = frame[top:bottom, left:right]
get_rect
函数根据 box
和 frame
返回四个坐标值(上、左、下、右),这些坐标用于定义要分析的矩形区域,并从 frame
中提取出这个子区域。
之后,对提取出的帧部分进行角度计算和修正:
python
angle_flow = correct_flow(frame_rect, model_seg, low_flow, fix_up_value, fix_down_value, debug_show == 1)
correct_flow
函数使用模型 model_seg
来预测并计算角度,同时利用 fix_up_value
和 fix_down_value
对结果进行了调整。
最后,检查 angle_flow
是否为 None
,如果是,则返回一个错误信息,表示流量识别失败。否则,返回修正后的流量值和一个空字符串作为成功的标志:
python
if angle_flow is None:
return None, f'识别失败:[{position}]预测失败!请重试或换表!'
return angle_flow, ''
完整代码:
python
def get_real_flow(box, frame, model_seg, low_flow, position, debug_show):
fix_up_value = '0,1,2,3'
fix_down_value = '7,8,9'
# x0.01 和 x0.1
top, left, bottom, right = get_rect(box, frame)
frame_rect = frame[top:bottom, left:right]
# 角度计算、修正
angle_flow = correct_flow(frame_rect, model_seg, low_flow, fix_up_value, fix_down_value, debug_show == 1)
if angle_flow is None:
return None, f'识别失败:[{position}]预测失败!请重试或换表!'
return angle_flow, ''
角度计算与读数修正:
python
def correct_flow(img, model_seg, low_flow, fix_up_value, fix_down_value, debug_show):
'''
角度计算与修正
:param img:
:param model_seg: 指针分割模型
:param low_flow: 低位值
:param fix_up_value: 需要向上修正的列表
:param fix_down_value: 需要向下修正的列表
:param debug_show:
:return: 角度对应的刻度值
'''
try:
circle_center, pointer_top = seg_center_top(img, model_seg, debug_show)
except:
return None
if circle_center is None or pointer_top is None:
return None
zero_top = (circle_center[0], 20)
# 求夹角
B = cal_ang(zero_top, circle_center, pointer_top)
if pointer_top[0] < circle_center[0]:
B = 360 - B
# 角度换算
per_scale = 3.6
# 向下取整 通过角度计算出来的
angle_flow = math.floor((B / 10) / per_scale)
# 修正算法.. 此处省略
return angle_flow
分割模型预测结果处理:
定义一个名为 seg_center_top
的函数,其主要目的是使用一个分割模型来预测图像中的指针和刻度圆的位置,并计算出刻度圆的中心点和指针的顶点位置。
函数签名
python
def seg_center_top(img, model_seg, debug_show):
img
: 输入的目标检测出来的指针区域图。model_seg
: 用于预测的分割模型。debug_show
: 调试模式标志位,当设置为 True 时,会显示一些中间结果以供调试。
图像预处理
python
img_copy = cv2.pyrUp(img.copy())
img_copy = cv2.addWeighted(img_copy, 1.2, img_copy, 0, 0)
- 使用
cv2.pyrUp
方法对输入图像进行上采样,以提高分辨率。 - 使用
cv2.addWeighted
增加图像亮度,这可能有助于改善分割模型的预测效果。
分割预测
python
results = model_seg.predict(img_copy, imgsz=128, verbose=False)
- 使用分割模型对处理后的图像进行预测,获取预测结果。
处理预测结果
接下来的部分代码处理了分割模型返回的结果,提取了两个类别的mask(指针和刻度圆),并确定哪个是刻度圆,哪个是指针。如果检测到的对象数量小于等于1则直接返回 None。
计算中心点和指针顶点
python
w = x2 - x1
h = y2 - y1
center = ((w / 2) + x1, (h / 2) + y1)
cx, cy = center[0], center[1]
- 计算了边界框的宽度和高度,并根据这些信息计算出刻度圆的中心点坐标
(cx, cy)
。
python
distances = np.sqrt(np.sum((pointers - np.array([cx, cy])) ** 2, axis=1))
top_pt = pointers[np.argmax(distances)]
- 计算了指针上的所有点与刻度圆中心的距离,并选择距离最远的点作为指针的顶点。
完整代码:
python
def seg_center_top(img, model_seg, debug_show):
'''
使用分割模型预测,并找到指针顶点和刻度圆中点
:param img: 目标检测出来的指针区域图
:param debug_show: 调试模式
:return:
'''
img_copy = cv2.pyrUp(img.copy())
# 增加亮度
img_copy = cv2.addWeighted(img_copy, 1.2, img_copy, 0, 0)
# 预测
results = model_seg.predict(img_copy, imgsz=128, verbose=False)
for result in results:
box_datas = result.boxes.data.cpu().numpy()
if len(box_datas) <= 1:
return None, None
p1 = result.masks.xy[0]
p2 = result.masks.xy[1]
# 一般情况下刻度圆circle_points的mask长度比指针pointers的mask要多
circle_points, pointers = (p2, p1) if len(p2) > len(p1) else (p1, p2)
x1, y1, x2, y2, conf, _ = box_datas[0]
for i, box_data in enumerate(box_datas):
cls = int(box_data[5])
if cls == 1:
# 指针刻度圆box
x1, y1, x2, y2, conf, _ = box_data
# 指针mask
pointers, circle_points = (result.masks.xy[1], result.masks.xy[0]) if i == 0 else (
result.masks.xy[0], result.masks.xy[1])
break
w = x2 - x1
h = y2 - y1
center = ((w / 2) + x1, (h / 2) + y1)
cx, cy = center[0], center[1]
# 求指针顶点
distances = np.sqrt(np.sum((pointers - np.array([cx, cy])) ** 2, axis=1))
top_pt = pointers[np.argmax(distances)]
if debug_show:
# 预测出来的box
cv2.rectangle(img_copy, (math.ceil(x1), math.ceil(y1)), (math.ceil(x2), math.ceil(y2)), (0, 0, 255),
thickness=2)
# 圆心
cv2.circle(img_copy, (math.ceil(cx), math.ceil(cy)), 2, (0, 255, 0), thickness=2)
# 指针顶点
cv2.circle(img_copy, (math.ceil(top_pt[0]), math.ceil(top_pt[1])), 2, (0, 255, 0), thickness=2)
# 连线圆心和指针顶点
cv2.line(img_copy, (math.ceil(cx), math.ceil(cy)), (math.ceil(top_pt[0]), math.ceil(top_pt[1])),
(0, 255, 0), thickness=2)
cv2.imshow('img', img_copy)
cv2.waitKey(0)
cv2.destroyAllWindows()
return (cx, cy), top_pt
字轮数字预测、修正、读数
由于篇幅较长,此处省略... 感兴趣的小伙伴可以私信!