pytorch yolov5 指针表计识别 分步识别表计
·

根据这三个矩形对角线的夹角求值
import math
import torch
from PIL import Image
import cv2 as cv
import matplotlib.pyplot as plt
import os
res = []
def addline(img, ls):
try:
l1 = ls.groupby('name').apply(lambda group: group[group['name'] == 'st']).values[0][0:4]
except:
return 0
try:
l2 = ls.groupby('name').apply(lambda group: group[group['name'] == 'end']).values[0][0:4]
except:
return 0
try:
l3 = ls.groupby('name').apply(lambda group: group[group['name'] == 'point']).values[0][0:4]
except:
return 0
cvimg = cv.imread(img)
x1 = round(l1[0],0)
x2 = round(l1[1],0)
y1 = round(l1[2],0)
y2 = round(l1[3],0)
print(x1)
cv.line(cvimg, (x1, y1), (x2, y2), (100, 100, 0), 2)
x1 = l2[0]
x2 = l2[1]
y1 = l2[2]
y2 = l2[3]
cv.line(cvimg, (x1, y1), (x2, y2), (150, 100, 0), 2)
x1 = l3[0]
x2 = l3[1]
y1 = l3[2]
y2 = l3[3]
cv.line(cvimg, (x1, y1), (x2, y2), (150, 200, 0), 2)
cv.imwrite(img, cvimg)
return 1
def line_k(l):
x1 = l[0]
x2 = l[1]
y1 = l[2]
y2 = l[3]
k = (y2 - y1) / (x2 - x1)
return k
def caldeg(ls):
# 分组筛选概率最大的三条
# ls = linelist.groupby('name').apply(lambda group: group[group['confidence'] == group['confidence'].max()])
l1 = ls.groupby('name').apply(lambda group: group[group['name'] == 'st']).values[0][0:4]
l2 = ls.groupby('name').apply(lambda group: group[group['name'] == 'end']).values[0][0:4]
l3 = ls.groupby('name').apply(lambda group: group[group['name'] == 'point']).values[0][0:4]
k1 = line_k(l1)
k2 = line_k(l2)
k3 = line_k(l3)
d1 = math.degrees(math.atan(k1))
d2 = math.degrees(math.atan(k2))
d3 = math.degrees(math.atan(k3))
degrees = (d3 - d1) / (d2 - d1) * 0.2 + 0.5
return degrees
class ModelRec():
def __int__(self):
self.img = ''
self.pt = ''
# self.weight = "det_helmet"
def imgrec(self):
self.model = torch.hub.load("./", "custom", path=self.pt,
source="local") # 加载安全帽检测模型
self.model.conf = 0.5
results = self.model(self.img)
# print(results)
return results
def img_cut(img1, rg, img2):
img = Image.open(img1)
# rg=(0, 0, 50, 50)
region = img.crop(rg)
region.save(img2)
if __name__ == '__main__':
img1 = './data/images/18.jpg' # 指定识别照片
img1_path=os.path.dirname(img1)
img1_name=os.path.basename(img1)
img1_name1=os.path.splitext(img1_name)
img2=img1_path + '/' + img1_name1[0] + 'cut' + img1_name1[1]
img3 = img1_path + '/' + img1_name1[0] + 'addline' + img1_name1[1]
pt1 = "./models/best_dis.pt"
pt2 = "./models/best_rec.pt"
model1 = ModelRec()
model1.pt = pt1
model2 = ModelRec()
model2.pt = pt2
model1.img = img1
result1 = model1.imgrec()
if result1.pandas().xyxy[0].__len__() > 0:
# --------- 截取图片 进入第二步识别- 多表计未考虑---
print(result1)
name = result1.pandas().xyxy[0].name.values[0]
img_rg = result1.pandas().xyxy[0].values[0][0:4]
img_cut(img1, img_rg, img2)
model2.img = img2
result2 = model2.imgrec()
ls=result2.pandas().xyxy[0]
if ls.__len__() > 0:
1
# 准备划线 --------------------
#addline(img3, ls)
if result2.pandas().xyxy[0].__len__() >= 3:
degrees = caldeg(ls)
print('识别的图片为', name, '压力为:', round(degrees, 2))
else:
print('关键点未识别')
else:
print('未识别出表计')

数据集和 训练模型下载地址 https://download.csdn.net/download/oMoZhe/88123412
更多推荐



所有评论(0)