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, 94, 161], 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 = cv2.KalmanFilter(4, 2)  # 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 update_color_threshold(color_name):
    global lower_hsv, upper_hsv
    # 根据选择的颜色更新颜色阈值
    lower_hsv = color_hsv_ranges[color_name]['lower']
    upper_hsv = color_hsv_ranges[color_name]['upper']


# 开始摄像头处理线程
def start_camera():
    global running
    if not running:
        running = True
        thread = Thread(target=camera_thread)
        thread.start()
        start_button.config(text="Stop Camera")
    else:
        running = False
        start_button.config(text="Start Camera")


# 摄像头处理线程函数
def camera_thread():
    global running, initial_measurement
    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 = np.ones((5, 5), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)

        # 查找轮廓
        contours, _ = cv2.findContours(mask.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 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)


# 更新自定义颜色的HSV阈值
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()


# 创建主窗口
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])

# 初始化摄像头
cap = cv2.VideoCapture(0)

# 启动GUI主循环
root.mainloop()