opencv/工程/Color_tracking.py

import cv2
import numpy as np
import tkinter as tk
from tkinter import colorchooser, simpledialog, Scale, HORIZONTAL, messagebox
from threading import Thread

# 颜色的HSV阈值
color_hsv_ranges = {
    'blue': {'lower': np.array([100, 150, 50], dtype=np.uint8),
             'upper': np.array([140, 255, 255], dtype=np.uint8)},
    'green': {'lower': np.array([35, 46, 120], dtype=np.uint8),
              'upper': np.array([60, 255, 255], dtype=np.uint8)},
    'red': {'lower': np.array([0, 116, 164], dtype=np.uint8),
            'upper': np.array([180, 255, 255], dtype=np.uint8)},
    'custom': {'lower': np.array([0, 0, 0], dtype=np.uint8),
               'upper': np.array([180, 255, 255], dtype=np.uint8)}
}

# 初始化选择的颜色
selected_color = 'green'
#获取初始颜色HSV阈值
lower_hsv, upper_hsv = (color_hsv_ranges[selected_color]['lower'],
                        color_hsv_ranges[selected_color]['upper'])

# 全局变量
running = False
kf = None
cap = None
root = None
def initialize_kalman_filter():
    """初始化卡尔曼滤波器"""
    global kf
    kf = cv2.KalmanFilter(4, 2)
    kf.measurementMatrix = np.array([[1, 0, 0, 0],
                                      [0, 1, 0, 0]], np.float32)
    kf.transitionMatrix = np.array([[1, 0, 1, 0],
                                    [0, 1, 0, 1],
                                    [0, 0, 1, 0],
                                    [0, 0, 0, 1]], np.float32)
    kf.processNoiseCov = np.array([[1, 0, 0, 0],
                                   [0, 1, 0, 0],
                                   [0, 0, 1, 0],
                                   [0, 0, 0, 1]], np.float32) * 0.03

# 用于初始化卡尔曼滤波器的变量
initial_measurement = None

def initialize_camera():
    """初始化摄像头"""
    global cap
    cap = cv2.VideoCapture(0)

def choose_color():
    global selected_color
    # 显示颜色选择对话框
    color_options = list(color_hsv_ranges.keys())
    selected_color = simpledialog.askstring("Choose Color", "Enter the color (blue, green, red, custom):",
                                            initialvalue=selected_color, parent=root)
    # 检查颜色选择是否有效，更新阈值
    if selected_color in color_options:
        update_color_threshold(selected_color)
    else:
        messagebox.showerror("Error", "Invalid color selection.")

def choose_custom_color():
    color_code = colorchooser.askcolor(title="Choose Custom Color")
    if color_code[0] is not None:
        lower = np.array([color_code[0][0] / 2, color_code[0][1] / 2, color_code[0][2] / 2])
        upper = np.array([color_code[0][0], color_code[0][1], color_code[0][2]])
        update_custom_hsv(lower, upper)

def update_custom_hsv(lower, upper):
    global color_hsv_ranges
    color_hsv_ranges['custom']['lower'] = lower.astype(np.uint8)
    color_hsv_ranges['custom']['upper'] = upper.astype(np.uint8)
    update_color_threshold('custom')

def exit_app():
    global running, cap
    running = False
    cap.release()
    root.quit()

def initialize_gui():
    """初始化GUI界面"""
    global root
    root = tk.Tk()
    root.title("Color Tracker")
    # 开始/停止摄像头按钮
    start_button = tk.Button(root, text="开始/停止", command=start_camera)
    start_button.pack()

    # 颜色选择按钮
    choose_color_button = tk.Button(root, text="选择颜色", command=choose_color)
    choose_color_button.pack()

    # 自定义颜色按钮
    custom_color_button = tk.Button(root, text="自定义颜色", command=choose_custom_color)
    custom_color_button.pack()

    # 退出按钮
    exit_button = tk.Button(root, text="退出", command=exit_app)
    exit_button.pack()

    # 自定义颜色HSV滑块
    def set_custom_hsv(lower_h, lower_s, lower_v, upper_h, upper_s, upper_v):
        update_custom_hsv(
            np.array([lower_h, lower_s, lower_v]),
            np.array([upper_h, upper_s, upper_v])
        )

    # 自定义颜色HSV滑块
    lower_h_scale = Scale(root, from_=0, to=180, orient=HORIZONTAL, label="Hue Min",
                          command=lambda v: set_custom_hsv(v, lower_s_scale.get(), lower_v_scale.get(),
                                                           upper_h_scale.get(),
                                                           upper_s_scale.get(), upper_v_scale.get()))
    lower_h_scale.pack()

    lower_s_scale = Scale(root, from_=0, to=255, orient=HORIZONTAL, label="Saturation Min",
                          command=lambda v: set_custom_hsv(lower_h_scale.get(), v, lower_v_scale.get(),
                                                           upper_h_scale.get(),
                                                           upper_s_scale.get(), upper_v_scale.get()))
    lower_s_scale.pack()

    lower_v_scale = Scale(root, from_=0, to=255, orient=HORIZONTAL, label="Value Min",
                          command=lambda v: set_custom_hsv(lower_h_scale.get(), lower_s_scale.get(), v,
                                                           upper_h_scale.get(),
                                                           upper_s_scale.get(), upper_v_scale.get()))
    lower_v_scale.pack()

    upper_h_scale = Scale(root, from_=0, to=180, orient=HORIZONTAL, label="Hue Max",
                          command=lambda v: set_custom_hsv(lower_h_scale.get(), lower_s_scale.get(),
                                                           lower_v_scale.get(), v,
                                                           upper_s_scale.get(), upper_v_scale.get()))
    upper_h_scale.pack()

    upper_s_scale = Scale(root, from_=0, to=255, orient=HORIZONTAL, label="Saturation Max",
                          command=lambda v: set_custom_hsv(lower_h_scale.get(), lower_s_scale.get(),
                                                           lower_v_scale.get(),
                                                           upper_h_scale.get(), v, upper_v_scale.get()))
    upper_s_scale.pack()

    upper_v_scale = Scale(root, from_=0, to=255, orient=HORIZONTAL, label="Value Max",
                          command=lambda v: set_custom_hsv(lower_h_scale.get(), lower_s_scale.get(),
                                                           lower_v_scale.get(),
                                                           upper_h_scale.get(), upper_s_scale.get(), v))
    upper_v_scale.pack()

    # 初始化滑块值
    lower_h_scale.set(color_hsv_ranges['custom']['lower'][0])
    lower_s_scale.set(color_hsv_ranges['custom']['lower'][1])
    lower_v_scale.set(color_hsv_ranges['custom']['lower'][2])
    upper_h_scale.set(color_hsv_ranges['custom']['upper'][0])
    upper_s_scale.set(color_hsv_ranges['custom']['upper'][1])
    upper_v_scale.set(color_hsv_ranges['custom']['upper'][2])

def start_camera():
    """开始或停止摄像头处理"""
    global running
    if not running:
        running = True
        thread = Thread(target=camera_thread)
        thread.start()
    else:
        running = False

def camera_thread():
    """摄像头处理线程函数"""
    global running, initial_measurement, lower_hsv, upper_hsv
    while running:
        ret, frame = cap.read()
        if not ret:
            messagebox.showerror("Error", "Failed to read frame from camera.")
            running = False
            break

        # 将BGR图像转换为HSV
        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
        mask = cv2.inRange(hsv, lower_hsv, upper_hsv)

        # 应用开运算（先腐蚀后膨胀）以去除噪声
        kernel_opening = np.ones((5, 5), np.uint8)
        mask_opened = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel_opening)

        # 应用闭运算（先膨胀后腐蚀）以填充孔洞
        kernel_closing = np.ones((5, 5), np.uint8)
        mask_closed = cv2.morphologyEx(mask_opened, cv2.MORPH_CLOSE, kernel_closing)

        # 使用闭合后的掩模找到轮廓
        contours, _ = cv2.findContours(mask_closed.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        # 遍历轮廓，画出边界框
        for contour in contours:
            if cv2.contourArea(contour) > 100:
                x, y, w, h = cv2.boundingRect(contour)
                cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)

                # 初始化卡尔曼滤波器的测量值
                if initial_measurement is None:
                    initial_measurement = np.array([[x], [y]], np.float32)
                    kf.statePost = initial_measurement
                else:
                    # 更新卡尔曼滤波器
                    measurement = np.array([[x], [y]], np.float32)
                    kf.correct(measurement)
                    prediction = kf.predict()
                    cv2.rectangle(frame, (int(prediction[0]), int(prediction[1])),
                                  (int(prediction[0] + w), int(prediction[1] + h)),
                                  (255, 0, 0), 2)

        cv2.imshow('Frame', frame)
        cv2.imshow('Mask', mask)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            running = False

    cv2.destroyAllWindows()

def update_color_threshold(color_name):
    """更新颜色阈值"""
    global lower_hsv, upper_hsv, selected_color
    # 根据选择的颜色更新颜色阈值
    selected_color = color_name  # 更新selected_color
    lower_hsv = color_hsv_ranges[color_name]['lower']
    upper_hsv = color_hsv_ranges[color_name]['upper']


def main():
    """主函数，程序入口"""
    initialize_kalman_filter()
    initialize_camera()
    initialize_gui()
    root.mainloop()

if __name__ == "__main__":
    main()