yolov5的二次开发以及目标缺陷检测项目（一）

一、解决的问题和目的

       最近做了一个小项目，主要是根据对目标物体的各个部位进行检测和识别，最终的目的是去判断目标物体对于区域位置是否有破损、烧毁等情况发生。

        本次拿电能表举例。目的想要检测屏幕、指示灯、键盘区域位置是否破损，屏幕亮暗是否正常，以及各指示灯的亮灭是否符合正常规律。

       

 图1

 二、实现步骤和方法

        首先通过yolov5对适量数据图片进行标注，我使用的labeImg进行的标注，然后经过yolov5项目进行数据集的训练。其中标注和训练步骤有很多资源可以借鉴，本次不再赘述。

提供以下参考链接：

labelmg标注：使用labelme打标签，详细教程_labelme使用教程-CSDN博客

yolov5模型训练：目标检测---教你利用yolov5训练自己的目标检测模型_路 blog . csdn . net / didiaopao / category _113216-CSDN博客        将模型训练好后，进行检测就得到图1效果。再进一步对每个检测内容做一些图像处理的操作，我将其定义为模型的后处理部分内容。

三、模型后处理

0、screen
1、runlight
2、warnlight
3、powerlight
4、button

这些是我要识别的对象，也是标注是的classes。至于为什么要把运行指示灯、警告指示灯、电源指示灯单拿出来标注，其中有一个小策略，在进行标注过程中可以将指示灯位置和指示灯下面的文字一起作为信息进行框选，这样就不用进一步对模型推理后的指示灯做区分了，索引为1的就是运行指示灯，2就是警告指示灯，3就是电源指示灯。

        (1)、针对屏幕的亮暗判断，可以通过yolov5模型推理出来的数据（中心点位置，类别，以及宽高）等信息，在原图上进行提取，然后进行裁剪为ROI（感兴趣区域）操作，最后再进行图像处理，可以通过输出的ROI图像的灰度平均值来对屏幕区域图像进行处理。设定一个阈值150（举例），其灰度平均值大于150，则为亮屏，反之则为暗屏。

方法一：

        根据yolov5输出的信息进行判断，在模型训练之前就对黑屏和亮屏设置为两个检测对象，这样就可以直接通过输出判断屏幕亮灭。

import cv2
import numpy as np
import matplotlib.pyplot as plt
import os
import sys
from terminal_data.DetectProject.Postprocessing.tools import Tools as rt
from terminal_data.DetectProject.Postprocessing.tools import roiRelatedOperations as ROI
from terminal_data.DetectProject.Postprocessing.tools import readloadPath
from terminal_data.DetectProject.Postprocessing.IndicatorLight import indicatorLightDetection as ILD


class calculate_MeanValue:
    def __init__(self, img):
        self.img = img
        self.hist = self.hist()
    def hist(self):
        plt.rcParams['font.family'] = 'SimHei'
        # 获取原图像的灰度直方图

        hist, bins = np.histogram(self.img.flatten(), 256, [0, 256])
        # # 为图像添加标题
        plt.title("灰度直方图")
        # 绘制灰度直方图
        plt.hist(self.img.flatten(), 256, [0, 256], color='gray')
        plt.xlim([0, 256])
        plt.xlabel('灰度值')
        plt.ylabel('像素数量')
        plt.show()
        return hist

    def calculateMean(self):
        # 计算概率密度函数
        pdf = self.hist / float(np.sum(self.hist))
        mean_value = np.sum(pdf * np.arange(256))
        return mean_value


if __name__ == "__main__":
    # 读取彩色图像并转换为灰度图像
    img_ori = cv2.imread("../../testDataset/260.jpg")
    txt1 = readloadPath.parseTxt("260.txt")
    path1 = txt1.getTxt()
    RD1 = rt.readTxt(path1, 0)
    values1 = RD1.open_file()
    if values1 == None:
        print("Input's index error, please check index have exist in txt file.")
        exit()
    img = ROI.extract_roi(img_ori, values1)
    cal = calculate_MeanValue(img)
    mean_value = cal.calculateMean()
    print("the picutre's mean_value is : ", mean_value)
    # gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # # 使用Otsu算法进行图像二值化
    # thre, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    # print("the picutre's threshold is : ", thre)
    #
    # # 显示原始图像和二值化图像
    # cv2.imshow('Original Image', img)
    # cv2.imshow('Binary Image', binary)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

方法二：

        根据yolov5输出并裁剪后的ROI图片，计算其灰度平均值方式判断屏幕图像是否亮或灭。

import cv2
import numpy as np
import matplotlib.pyplot as plt
import os
import sys
from terminal_data.DetectProject.Postprocessing.tools import Tools as rt
from terminal_data.DetectProject.Postprocessing.tools import roiRelatedOperations as ROI
from terminal_data.DetectProject.Postprocessing.tools import readloadPath
from terminal_data.DetectProject.Postprocessing.IndicatorLight import indicatorLightDetection as ILD


class calculate_MeanValue:
    def __init__(self, img):
        self.img = img
        self.hist = self.hist()
    def hist(self):
        plt.rcParams['font.family'] = 'SimHei'
        # 获取原图像的灰度直方图

        hist, bins = np.histogram(self.img.flatten(), 256, [0, 256])
        # # 为图像添加标题
        plt.title("灰度直方图")
        # 绘制灰度直方图
        plt.hist(self.img.flatten(), 256, [0, 256], color='gray')
        plt.xlim([0, 256])
        plt.xlabel('灰度值')
        plt.ylabel('像素数量')
        plt.show()
        return hist

    def calculateMean(self):
        # 计算概率密度函数
        pdf = self.hist / float(np.sum(self.hist))
        mean_value = np.sum(pdf * np.arange(256))
        return mean_value


if __name__ == "__main__":
    # 读取彩色图像并转换为灰度图像
    img_ori = cv2.imread("../../testDataset/260.jpg")
    txt1 = readloadPath.parseTxt("260.txt")
    path1 = txt1.getTxt()
    RD1 = rt.readTxt(path1, 0)
    values1 = RD1.open_file()
    if values1 == None:
        print("Input's index error, please check index have exist in txt file.")
        exit()
    img = ROI.extract_roi(img_ori, values1)
    cal = calculate_MeanValue(img)
    mean_value = cal.calculateMean()
    print("the picutre's mean_value is : ", mean_value)
    # gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # # 使用Otsu算法进行图像二值化
    # thre, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    # print("the picutre's threshold is : ", thre)
    #
    # # 显示原始图像和二值化图像
    # cv2.imshow('Original Image', img)
    # cv2.imshow('Binary Image', binary)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

        (2)、对于键盘区域是否完整，判断其是否存在缺陷，可以建立一个标准且完整的键盘图片作为基准，将yolov5模型推理并经过裁剪的ROI图像与基准图片进行对比，根据其相似性方法，根据检测图像与基准图像相似度进行判断，从而获得检测结果，判断键盘是否完整。

键盘完整度检测

import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import matplotlib.pyplot as plt
import os
import sys
from terminal_data.DetectProject.Postprocessing.tools import Tools as rt
from terminal_data.DetectProject.Postprocessing.tools import roiRelatedOperations as ROI
from terminal_data.DetectProject.Postprocessing.tools import readloadPath


# 可以准确识别按键完整度
#用于给图片添加中文字符的函数
def cv2ImgAddText(img, text, left, top, textColor=(0, 255, 0), textSize=20):
    # 判断是否OpenCV图片类型
    if (isinstance(img, np.ndarray)):
        img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    # 创建一个可以在给定图像上绘图的对象
    draw = ImageDraw.Draw(img)
    # 字体的格式
    fontStyle=ImageFont.truetype("font/simsun.ttc",textSize, encoding="utf-8")
    # 绘制文本
    draw.text((left, top), text, textColor, font=fontStyle)
    # 转换回OpenCV格式
    return cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)

if __name__ == "__main__":
    # plt绘图显示中文
    plt.rcParams['font.family'] = 'SimHei'

    img0 = cv2.imread("../../testDataset/260.jpg")
    txt1 = readloadPath.parseTxt("260.txt")
    path1 = txt1.getTxt()
    RD1 = rt.readTxt(path1, 4)
    values1 = RD1.open_file()
    img_roi1 = ROI.extract_roi(img0, values1)
    cv2.imshow('img_roi对照图', img_roi1)
    # 彩色图转灰度图
    img1 = cv2.cvtColor(img_roi1, cv2.COLOR_BGR2GRAY)
    # 缺陷检测图
    img = cv2.imread("../../testDataset/68.jpg")
    txt = readloadPath.parseTxt("68.txt")
    path = txt.getTxt()
    RD = rt.readTxt(path, 4)
    values = RD.open_file()
    img_roi = ROI.extract_roi(img, values)
    cv2.imshow(rt.zh_ch("roi实验图"), img_roi)
    # 获取灰度图像
    defect_img1 = cv2.cvtColor(img_roi, cv2.COLOR_BGR2GRAY)


    # 获取原图像的灰度直方图
    hist0 = cv2.calcHist([img1], [0], None, [256], [0.0, 255.0])
    # 获取待检测图像的灰度直方图
    hist1 = cv2.calcHist([defect_img1], [0], None, [256], [0.0, 255.0])
    # 为图像添加标题
    plt.title("原图与待检测img对比")
    # 添加图例
    plt.plot(hist0, label='原图')
    plt.plot(hist1, label='待检测img')
    # 相似度比较
    rst = cv2.compareHist(hist0, hist1, method=cv2.HISTCMP_CORREL)
    print("Similarity is : ", rst)
    # res >= 0.95即认为合格
    # defect_img0.resize(416, 416)
    cv2.imshow(rt.zh_ch("img结果图"), cv2ImgAddText(img_roi, "合格" if rst >= 0.3 else "不合格", 20, 20, (255, 0, 0), 25))
    # 设置x轴的数值范围
    plt.xlim([0, 256])
    plt.legend(loc='upper left')
    plt.show()
    cv2.waitKey(0)

        (3)、对于指示灯的亮灭，可以和(1)类似，判断裁剪出来的ROI的图像亮度平均值，来判断指示灯状态。然后根据实际运行逻辑来判断整个电能表运行过程总各指示灯是否运行正常。
# 指示灯亮灭检测

import cv2 as cv
import numpy as np
import os
import matplotlib.pyplot as plt
from terminal_data.DetectProject.Postprocessing.tools import Tools as rt
from terminal_data.DetectProject.Postprocessing.tools import readloadPath
from terminal_data.DetectProject.Postprocessing.tools import roiRelatedOperations as ROI

# # 获取当前目录的绝对路径
# current_dir = os.getcwd()
# # 设置当前目录为根目录
# os.chdir(current_dir)

# 指示灯亮灭检测
# 指示灯亮灭判断
class pilotLampStatus():
    def __init__(self, img, direction = 0): # d=direction为0则处理图像左边部分， 为0则处理右边部分
        self.img = img
        self.h, self.w, self.c = self.img.shape
        self.x, self.y, self.r = None, None, None
        self.list = []       # 存一个二维数组
        self.direction = direction
        self.hist = self.getHistogram()

    # 图像预处理（灰度化、高斯滤波、检测圆形区域）
    def preprocess(self):
        gray = cv.cvtColor(self.img, cv.COLOR_BGR2GRAY)
        # gray1 = cv.resize(gray, (200, 200))  # resize为：w*h
        # cv.imshow("gray", gray1)
        # 将图像高斯滤波降噪处理
        blur = cv.GaussianBlur(gray, (5, 5), 0)
        # 使用 Hough 变换检测圆形   霍夫圆变换算法检测
        circles = cv.HoughCircles(
            blur,
            cv.HOUGH_GRADIENT,
            dp=1,
            minDist=20,
            param1=50,
            param2=10,    # 累加器阈值，越小则检测到的圆更多但会有误检测
            minRadius=0,
            maxRadius=0
        )
        # 如果有检测到圆形
        if circles is not None:
            # 将结果转换为整数
            circles = np.round(circles[0, :]).astype("int")
            # print("size is :", circles.size)
            # 在图像上绘制圆形
            for (x, y, r) in circles:
                list = []
                # cv.circle(self.img, (x, y), r, (0, 255, 0), 2)
                list.append(x)
                list.append(y)
                list.append(r)
                self.list.append(list)
            # self.img = cv.resize(self.img, (200, 200))
            # cv.imshow("gray", self.img)
            return self.list
        else:
            print("don't detect circles in pic!")
            return

    # 提取图像中的圆形
    def extractCircle(self):          # d=direction为0则处理图像左边部分， 为0则处理右边部分
        list = self.preprocess()
        if len(list) == 0:
            print("Can not find circle in picture!")
            return
        #选择运行还是告警指示灯：0为左边图运行指示灯， 1为右边图告警指示灯
        list_son = list[self.direction]
        x, y, r = [int(list_son[i]) for i in range(len(list_son))]
        # 提取圆形区域内容
        extracted_region = self.img[y - r:y + r, x - r:x + r]
        #
        # # 显示提取的圆形区域内容
        # extracted_region = cv.resize(extracted_region, (200, 200))
        # cv.imshow("Extracted Region", extracted_region)
        return extracted_region

    # 左侧部分指示灯处理
    def left_process(self):
        img_left = self.img[:, 0:int(self.w/2)]
        cv.imshow("left", img_left)
        return img_left

    # 右侧部分指示灯处理
    def right_process(self):
        img_right = self.img[:, int(self.w/2):int(self.w)]
        cv.imshow("right", img_right)
        return img_right

    # 得到灰度直方图
    def getHistogram(self):
        plt.rcParams['font.family'] = 'SimHei'
        # 获取原图像的灰度直方图

        hist, bins = np.histogram(self.extractCircle().flatten(), 256, [0, 256])
        # # 为图像添加标题
        # plt.title("灰度直方图")
        # 绘制灰度直方图
        plt.hist(self.extractCircle().flatten(), 256, [0, 256], color='gray')
        # plt.xlim([0, 256])
        # plt.xlabel('灰度值')
        # plt.ylabel('像素数量')
        # plt.show()
        return hist

    # 计算灰度直方图的均值
    def calculateMean(self):
        hist = self.getHistogram()
        # 计算概率密度函数
        pdf = hist / float(np.sum(hist))
        mean_value = np.sum(pdf * np.arange(256))
        return mean_value

    # 解析最终的结果是亮还是暗
    def judgeResult(self):
        res = self.calculateMean()
        # print("res:", res)
        if res > 130:       # 阈值调整策略
            print("--------->result is bright!")
            return True
        else:
            print("--------->result is dark!")
            return False


if __name__ == "__main__":
    #58/60/68/156/250/259/260
    img_ori = cv.imread("../../testDataset/156.jpg")
    txt = readloadPath.parseTxt("156.txt")
    list_flag = []                           #其中存放运行灯、告警灯、电源指示灯的标志位， true为亮， false为灭
    # 先判断电源指示灯是否亮灭
    path = txt.getTxt()
    for i in range(1, 4):
        RD = rt.readTxt(path, i)            #其中参数1\2\3分别代表运行灯、告警灯、电源指示灯
        values = RD.open_file()
        if values == None:
            print("Input's index error, please check index have exist in txt file.")
            exit()
        img = ROI.extract_roi(img_ori, values)
        img = cv.resize(img, (200, 200))
        cv.imshow("img", img)
        cv.waitKey(0)
        # 实例化指示灯状态类
        PLS = pilotLampStatus(img)
        list_flag.append(PLS.judgeResult())
    print("list flag is : ", list_flag)
    ###################################判断灯是否正常################################################
    if list_flag[2] == False:
        if list_flag[0] == False and list_flag[1] == False:
            print("************Indicator light is normal!********************")
        elif list_flag[0] == True or list_flag[1] == True:
            print("-----------run light or warning light is error!-----------")
    else:
        if list_flag[0] == True and list_flag[1] == False:
            print("************Indicator light is normal!********************")
        elif list_flag[0] == False and list_flag[1] == True:
            print("-----------run light or warning light is error!-----------")
        else:
            print("-----------run light or warning light is error!-----------")

以上代码仅供参考，有一些逻辑暂时还未修改，项目具体思路和完整代码在下一篇博客中会发出来。后续还会将该项目整成http服务，作为服务端进行调用，客户端直接通道调用相应接口就可以调用输出检测结果。