基于OpenCV的ERD图形检测异常：矩形与菱形识别问题排查

问题场景 #

需要检测ERD（实体关系图）结构中的两个相连矩形和中间的菱形，但现有OpenCV代码未能正确识别两个矩形，未在矩形上标注Rectangle。

输入图像 #

当前输出图像 #

现有代码 #

import cv2
import numpy as np
from matplotlib import pyplot as plt

# 读取图像
img = cv2.imread('pic4.jpeg')

# 将图像转换为灰度图
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

# 设置灰度图阈值
_, threshold = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)

# 使用findContours()函数
contours, _ = cv2.findContours(
    threshold, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

i = 0

# 用于存储形状名称的列表
for contour in contours:

    # 忽略第一个轮廓，因为findContours函数会将整个图像识别为一个形状
    if i == 0:
        i = 1
        continue

    # 使用cv2.approxPolyDP()函数进行形状近似
    approx = cv2.approxPolyDP(
        contour, 0.01 * cv2.arcLength(contour, True), True)
    
    # 使用drawContours()函数绘制轮廓
    cv2.drawContours(img, [contour], 0, (0, 0, 0), 5)

    # 寻找形状的中心点
    M = cv2.moments(contour)
    if M['m00'] != 0.0:
        x = int(M['m10']/M['m00'])
        y = int(M['m01']/M['m00'])

    # 在每个形状的中心点标注形状名称
    if len(approx) == 3:
        cv2.putText(img, 'Triangle', (x, y),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2)

    elif len(approx) == 4:
        cv2.putText(img, 'Rectangle', (x, y),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2)

    elif len(approx) == 5:
        cv2.putText(img, 'Pentagon', (x, y),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2)

    elif len(approx) == 6:
        cv2.putText(img, 'Hexagon', (x, y),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2)

    else:
        cv2.putText(img, 'circle', (x, y),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2)

# 显示绘制轮廓后的图像
cv2.imshow('pic4', img)

cv2.waitKey(0)
cv2.destroyAllWindows()

问题原因 #

1. 阈值处理不合理：固定阈值127无法适配图像的亮度差异，导致二值化后矩形轮廓被干扰或合并。
2. 轮廓检索逻辑错误：cv2.RETR_TREE会检索所有层级轮廓，加上武断忽略第一个轮廓的逻辑，可能丢失有效矩形轮廓。
3. 形状近似参数不当：0.01的近似系数过小，导致矩形的近似多边形顶点数偏离4，无法被识别。

解决方案 #

优化方向 #

1. 预处理流程升级：增加模糊降噪，使用自适应阈值替代固定阈值，提升二值化效果。
2. 轮廓筛选逻辑调整：使用cv2.RETR_EXTERNAL只检索最外层轮廓，通过轮廓面积过滤无效小轮廓，取消忽略第一个轮廓的错误逻辑。
3. 形状判断精细化：调整近似系数，同时对4顶点形状通过邻边夹角区分矩形和菱形（矩形内角接近90度）。

修改后的代码 #

import cv2
import numpy as np

# 读取图像
img = cv2.imread('pic4.jpeg')
if img is None:
    print("无法读取图像，请检查路径")
    exit()

# 预处理：模糊降噪+自适应阈值
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
blurred = cv2.GaussianBlur(gray, (5, 5), 0)
threshold = cv2.adaptiveThreshold(blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 11, 2)

# 检索最外层轮廓
contours, _ = cv2.findContours(threshold, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

# 过滤过小的无效轮廓（可根据图像尺寸调整阈值）
min_area = 500
valid_contours = [cnt for cnt in contours if cv2.contourArea(cnt) > min_area]

for contour in valid_contours:
    # 轮廓近似
    epsilon = 0.02 * cv2.arcLength(contour, True)
    approx = cv2.approxPolyDP(contour, epsilon, True)
    
    # 绘制轮廓
    cv2.drawContours(img, [contour], 0, (0, 0, 0), 2)
    
    # 计算形状中心点
    M = cv2.moments(contour)
    if M['m00'] == 0.0:
        continue
    x = int(M['m10'] / M['m00'])
    y = int(M['m01'] / M['m00'])
    
    # 形状判断
    if len(approx) == 4:
        # 计算邻边向量夹角，区分矩形和菱形
        pts = approx.reshape(4, 2)
        v1 = np.subtract(pts[1], pts[0])
        v2 = np.subtract(pts[2], pts[1])
        cos_angle = abs(np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2)))
        
        if cos_angle < 0.1:  # 余弦值接近0，夹角接近90度
            cv2.putText(img, 'Rectangle', (x, y), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2)
        else:
            cv2.putText(img, 'Diamond', (x, y), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
    elif len(approx) == 3:
        cv2.putText(img, 'Triangle', (x, y), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2)
    elif len(approx) == 5:
        cv2.putText(img, 'Pentagon', (x, y), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2)
    else:
        cv2.putText(img, 'Circle', (x, y), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2)

# 显示结果
cv2.imshow('Shape Detection', img)
cv2.waitKey(0)
cv2.destroyAllWindows()

内容的提问来源于stack exchange，提问作者M.Umer Haider