气泡检测

气泡检测

# 胶囊瑕疵检测
import cv2
import os
import numpy as np


def bub_check(img_path, img_file, im, im_gray):
    # 二值化
    t, im_bin = cv2.threshold(im_gray,
                              170, 255,
                              cv2.THRESH_BINARY)
    cv2.imshow("im_bin", im_bin)

    # 查找轮廓
    img, cnts, hie = cv2.findContours(
        im_bin,  # 图像
        cv2.RETR_LIST,  # 检测所有轮廓
        cv2.CHAIN_APPROX_NONE)  # 存储所有轮廓点
    # 轮廓过滤
    new_cnts = []  # 过滤后的轮廓
    for cnt in cnts:
        area = cv2.contourArea(cnt)  # 计算轮廓面积
        print("area:", area)
        if area < 10000 and area > 22:
            new_cnts.append(cnt)  # 放入新的列表

    im_cnt = cv2.drawContours(im,  # 原图
                              new_cnts,  # 过滤后的轮廓数据
                              -1,  # 绘制所有轮廓
                              (0, 0, 255), 2)  # 颜色、粗细
    cv2.imshow("im_cnt", im_cnt)

    if len(new_cnts) > 0:
        print("气泡瑕疵：", img_path)
        # 可以将这些图像移动到另外的目录, 代码自己添加...


def balance_check(img_path, img_file, im, im_gray):
    # 模糊、膨胀，去掉过细细节
    im_blur = cv2.GaussianBlur(im_gray, (5, 5), 0)
    kernel = np.ones((5, 5), np.uint8)  # 计算核
    im_dilate = cv2.dilate(im_blur, kernel)
    cv2.imshow("im_dilate", im_dilate)

    # 边沿提取
    im_canny = cv2.Canny(im_dilate, 60, 200)
    cv2.imshow("im_canny", im_canny)

    # 查找、绘制轮廓
    img, cnts, hie = cv2.findContours(
        im_canny,
        cv2.RETR_LIST,
        cv2.CHAIN_APPROX_NONE)

    # 计算周长
    new_cnts = []  # 过滤后的轮廓
    if len(cnts) > 0:
        for c in cnts:  # 遍历每个轮廓
            circle_len = cv2.arcLength(c,  # 轮廓
                                       True)  # 是否封闭
            if circle_len >= 1000:  # 将大于1000的轮廓放入新列表
                new_cnts.append(c)

    # 按照面积排序
    new_cnts = sorted(new_cnts,  # 可迭代对象
                      key=cv2.contourArea,  # 排序依据
                      reverse=True)  # 倒序排列
    new_cnts = new_cnts[1:2]  # 取出面积第二大的轮廓

    im_cnt = cv2.drawContours(im,
                              new_cnts,
                              -1,
                              (0, 0, 255), 2)
    cv2.imshow("im_cnt", im_cnt)

    # 遍历轮廓，取出最大/小的x坐标，最大/小的y坐标
    max_x = new_cnts[0][0][0][0]
    max_y = new_cnts[0][0][0][1]
    min_x = max_x
    min_y = max_y

    for cnt in new_cnts[0]:
        if cnt[0][0] >= max_x:
            max_x = cnt[0][0]
        if cnt[0][0] <= min_x:
            min_x = cnt[0][0]
        if cnt[0][1] >= max_y:
            max_y = cnt[0][1]
        if cnt[0][1] <= min_y:
            min_y = cnt[0][1]
    # 绘制参考线
    # cv2.line(im, (min_x, min_y), (min_x, max_y), (0,0,255),2)
    # cv2.line(im, (min_x, min_y), (max_x, min_y), (0,0,255),2)
    # cv2.line(im, (max_x, min_y), (max_x, max_y), (0,0,255),2)
    # cv2.line(im, (min_x, max_y), (max_x, max_y), (0, 0, 255), 2)

    # 绘制水平中线
    center_y = int((max_y + min_y) / 2)
    center_up = int((center_y + min_y) / 2)  # 上中线
    center_down = int((center_y + max_y) / 2)  # 下中线

    # cv2.line(im, (min_x, center_y), (max_x, center_y),
    #          (0, 0, 255), 2)
    cv2.line(im, (min_x, center_up), (max_x, center_up),
             (0, 0, 255), 2)
    cv2.line(im, (min_x, center_down), (max_x, center_down),
             (0, 0, 255), 2)

    cv2.imshow("im_line", im)

    # 求药丸轮廓和上下中线交点
    cross_point_up = set()
    cross_point_down = set()

    for cnt in new_cnts[0]:
        x, y = cnt[0][0], cnt[0][1]  # 取出轮廓点的x,y坐标
        if y == center_up:  # 等于上中线y值
            cross_point_up.add((x, y))
        if y == center_down:  # 等于下中线y值
            cross_point_down.add((x, y))
    # 集合转列表
    cross_point_up = list(cross_point_up)
    cross_point_down = list(cross_point_down)

    # 标记交点
    for p in cross_point_up:
        cv2.circle(im,  # 原图
                   (p[0], p[1]), 8,  # 圆心,半径
                   (0, 0, 255), 2)
    for p in cross_point_down:
        cv2.circle(im,  # 原图
                   (p[0], p[1]), 8,  # 圆心,半径
                   (0, 0, 255), 2)
    cv2.imshow("im_circle", im)

    # 求上中线、下中线长度
    len_up, len_down = 0, 0
    p1 = cross_point_up[0]  # 上中线第一个点
    p2 = cross_point_up[1]  # 上中线第二个点
    len_up = abs(p1[0] - p2[0])  # 上中线长度

    p3 = cross_point_down[0]  # 下中线第一个点
    p4 = cross_point_down[1]  # 下中线第一个点
    len_down = abs(p3[0] - p4[0])  # 下中线长度

    print("len_up:", len_up)
    print("len_down:", len_down)

    if abs(len_up - len_down) >= 10:
        print("大小头:", img_path)
    else:
        print("正常大小")


if __name__ == "__main__":
    # 读取所有待检测图像
    img_dir = "capsules"
    img_files = os.listdir(img_dir)

    for img_file in img_files:
        # 拼接完整路径
        img_path = os.path.join(img_dir, img_file)
        if os.path.isdir(img_path):  # 是目录
            continue

        # 读取图像
        im = cv2.imread(img_path)
        im_gray = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
        cv2.imshow("im", im)
        cv2.imshow("im_gray", im_gray)

        # 气泡检测
        # bub_check(img_path, img_file, im, im_gray)

        # 大小头检测
        balance_check(img_path, img_file, im, im_gray)

        cv2.waitKey()
        cv2.destroyAllWindows()