根据这三个矩形对角线的夹角求值

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

更多推荐