基于相似性图识别带不完整紫色轮廓的Bot并获取Bounding Box参数

问题背景 #

需要为图像中带有紫色轮廓的每个Bot绘制Bounding Box，目前通过get_similarity_map函数生成的相似性图存在轮廓断点，无法完整追踪Bot个体，需实现Bot个体识别并获取其x、y坐标及宽度、高度参数。

解决方案 #

通过形态学操作修复轮廓断点，再结合OpenCV的轮廓检测与Bounding Box计算来实现需求，具体步骤如下：

• 对相似性图执行形态学闭运算，填补轮廓断点，连通Bot的完整轮廓
• 提取所有连通轮廓，过滤掉噪声小区域
• 计算每个有效轮廓的Bounding Box参数，并转换回原图坐标系
• 绘制Bounding Box并输出参数

修改后的完整代码 #

import cv2
import numpy as np
 
def get_similarity_map(image_hsv, target_hsv, hue_threshold, similarity_threshold):
    delta = np.abs(image_hsv - target_hsv)
    delta[:, :, 0] = np.minimum(delta[:, :, 0], 180 - delta[:, :, 0]) / 180
    delta[:, :, 1] /= 255
    delta[:, :, 2] /= 255
 
    similarity = 1 - np.sqrt(np.sum(delta[:, :, :] ** 2, axis=2)) / np.sqrt(3)
    mask = np.logical_and(
        delta[:, :, 0] <= hue_threshold,
        np.logical_and(
            similarity_threshold[0] <= similarity, similarity <= similarity_threshold[1]
        ),
    )
    return np.where(mask, similarity, 0)
 

purple = (150, 119, 179)
hue_threshold = 0.05
similarity_threshold = (0.75, 0.85)
crop_size = 1080
 
frame = cv2.imread("image.png")
 
height, width = frame.shape[:2]
roi_x = width // 2 - crop_size // 2
roi_y = height // 2 - crop_size // 2
roi = frame[roi_y : roi_y + crop_size, roi_x : roi_x + crop_size]
 
roi = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV)
roi = np.array(roi, dtype=np.float32)
 
similarity_map = get_similarity_map(roi, purple, hue_threshold, similarity_threshold)

# --- 新增处理步骤：修复轮廓断点 + 识别Bot个体 ---
# 转换为8位灰度图
similarity_map = (similarity_map * 255).astype(np.uint8)

# 形态学闭运算：先膨胀再腐蚀，填补断点
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
closed_map = cv2.morphologyEx(similarity_map, cv2.MORPH_CLOSE, kernel)

# 二值化，确保轮廓提取更准确
_, binary_map = cv2.threshold(closed_map, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

# 提取轮廓
contours, _ = cv2.findContours(binary_map, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

# 遍历轮廓，过滤噪声并计算Bounding Box
bot_list = []
for cnt in contours:
    # 过滤过小的噪声区域（可根据实际Bot大小调整阈值）
    if cv2.contourArea(cnt) < 50:
        continue
    # 计算ROI内的Bounding Box
    x_roi, y_roi, w, h = cv2.boundingRect(cnt)
    # 转换回原图的坐标
    x = roi_x + x_roi
    y = roi_y + y_roi
    bot_list.append({"x": x, "y": y, "width": w, "height": h})
    # 在原图上绘制Bounding Box
    cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)

# 输出每个Bot的参数
print("识别到的Bot参数：")
for idx, bot in enumerate(bot_list):
    print(f"Bot {idx+1}: x={bot['x']}, y={bot['y']}, width={bot['width']}, height={bot['height']}")

# 显示结果
cv2.imshow("Bot Detection Result", frame)
cv2.imwrite("bot_detection_result.png", frame)
cv2.waitKey(0)
cv2.destroyAllWindows()

关键步骤说明 #

1. 形态学闭运算：使用3x3的矩形结构元素，通过先膨胀后腐蚀的操作，填补相似性图中的小断点，让Bot的轮廓变得连通完整。
2. 二值化处理：通过OTSU自动阈值二值化，将相似性图转换为黑白二值图，减少轮廓提取的干扰。
3. 轮廓过滤：通过cv2.contourArea过滤掉面积过小的噪声区域，只保留Bot的有效轮廓。
4. 坐标转换：因为之前对原图做了ROI裁剪，所以需要将ROI内的Bounding Box坐标转换回原图的坐标系，保证绘制位置准确。

内容的提问来源于stack exchange，提问作者Jake Jin