📚 课程目录第2章基于颜色特征的工件分拣

2. 基于颜色特征的工件分拣

任务目标

本案例实现一个基于颜色直方图的工件分拣系统：通过摄像头采集传送带上的工件图像，提取HSV颜色直方图特征，使用预定义的参考颜色模板进行匹配，识别工件颜色后输出对应的分拣指令（料箱编号），模拟机械臂分拣控制信号。

环境与依赖

Python 3.8+
OpenCV（pip install opencv-python）
NumPy（pip install numpy）

核心实现

系统由三个模块组成：颜色直方图特征提取、参考颜色模板库、最近邻匹配分拣。HSV颜色空间对光照变化更鲁棒，因此将BGR图像转换到HSV后，对H（色相）和S（饱和度）通道计算2D直方图作为颜色特征。匹配时使用OpenCV的cv2.compareHist计算直方图相关性，选择相关性最高的颜色模板作为识别结果，并映射到对应料箱。

import cv2
import numpy as np

# ---------- 1. 颜色模板库（HSV直方图） ----------
# 每种颜色定义HSV范围，用于生成参考直方图
COLOR_DEFS = {
    'red':   {'hsv_range': ((0, 50, 50), (10, 255, 255)), 'bin': 1},
    'green': {'hsv_range': ((40, 50, 50), (80, 255, 255)), 'bin': 2},
    'blue':  {'hsv_range': ((100, 50, 50), (130, 255, 255)), 'bin': 3},
    'yellow': {'hsv_range': ((20, 50, 50), (35, 255, 255)), 'bin': 4},
}

def build_reference_hist():
    """为每种颜色生成标准HSV直方图（H:0-180, S:0-256）"""
    ref_hists = {}
    for name, cfg in COLOR_DEFS.items():
        # 创建一个纯色图像作为参考
        lower, upper = cfg['hsv_range']
        # 取HSV范围中间值作为代表色
        h = (lower[0] + upper[0]) // 2
        s = (lower[1] + upper[1]) // 2
        v = 255
        ref_img = np.zeros((50, 50, 3), dtype=np.uint8)
        ref_img[:, :] = (h, s, v)
        ref_img = cv2.cvtColor(ref_img, cv2.COLOR_HSV2BGR)
        # 转回HSV并计算直方图
        hsv = cv2.cvtColor(ref_img, cv2.COLOR_BGR2HSV)
        hist = cv2.calcHist([hsv], [0, 1], None, [30, 32],
                            [0, 180, 0, 256])
        cv2.normalize(hist, hist, 0, 1, cv2.NORM_MINMAX)
        ref_hists[name] = hist
    return ref_hists

# ---------- 2. 工件颜色识别 ----------
def classify_workpiece(bgr_img, ref_hists):
    """输入工件BGR图像，返回(颜色名称, 料箱编号)"""
    hsv = cv2.cvtColor(bgr_img, cv2.COLOR_BGR2HSV)
    # 计算待识别图像的HSV直方图
    hist = cv2.calcHist([hsv], [0, 1], None, [30, 32],
                        [0, 180, 0, 256])
    cv2.normalize(hist, hist, 0, 1, cv2.NORM_MINMAX)

    best_color = None
    best_score = -1
    for name, ref_hist in ref_hists.items():
        score = cv2.compareHist(hist, ref_hist, cv2.HISTCMP_CORREL)
        if score > best_score:
            best_score = score
            best_color = name
    return best_color, COLOR_DEFS[best_color]['bin']

# ---------- 3. 模拟分拣控制 ----------
def send_sort_signal(bin_id):
    """输出分拣指令（实际项目中此处通过串口/IO控制机械臂）"""
    print(f"[SORT] 工件分拣至料箱 #{bin_id}")

# ---------- 4. 主流程 ----------
def main():
    # 构建参考颜色直方图
    ref_hists = build_reference_hist()

    # 模拟3个待分拣工件（实际项目中由相机采集）
    test_images = {
        'sample_red.jpg':    np.full((100, 100, 3), (0, 0, 255), np.uint8),   # 红色
        'sample_green.jpg':  np.full((100, 100, 3), (0, 255, 0), np.uint8),   # 绿色
        'sample_blue.jpg':   np.full((100, 100, 3), (255, 0, 0), np.uint8),   # 蓝色
    }

    for fname, img in test_images.items():
        color, bin_id = classify_workpiece(img, ref_hists)
        print(f"{fname} -> 识别颜色: {color}")
        send_sort_signal(bin_id)

if __name__ == '__main__':
    main()

运行与验证

将上述代码保存为color_sort.py，在终端执行python color_sort.py。
程序使用纯色模拟图像验证分类逻辑。预期输出：

sample_red.jpg -> 识别颜色: red
[SORT] 工件分拣至料箱 #1
sample_green.jpg -> 识别颜色: green
[SORT] 工件分拣至料箱 #2
sample_blue.jpg -> 识别颜色: blue
[SORT] 工件分拣至料箱 #3

通过条件：三种颜色的识别结果与料箱编号均正确对应。若识别错误，检查HSV颜色范围定义是否覆盖目标颜色。

关键说明

HSV范围定义中，红色在OpenCV中H通道范围为0-10和170-180两段，本案例简化取0-10段；实际多红色场景需合并两段范围。
直方图通道数（30×32）和匹配方法（相关性）可根据实际工件颜色差异调整。颜色相近时增大直方图维度或改用cv2.HISTCMP_BHATTACHARYYA（巴氏距离，越小越相似）。
真实项目中，将main()中的模拟图像替换为相机实时帧，并将send_sort_signal()改为通过串口发送指令给机械臂控制器即可。