import cv2
import numpy as np

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 main():
    global initial_measurement, kf
    # 初始化摄像头
    cap = cv2.VideoCapture(0)

    # 用于标记是否已经选择颜色区域的变量
    color_selected = False
    color_lower = None
    color_upper = None

    while True:
        # 从摄像头读取一帧
        ret, frame = cap.read()
        if not ret:
            break

        # 如果还没有选择颜色区域，显示当前帧并允许用户选择
        if not color_selected:
            cv2.imshow('Select color area', frame)

            # 用户按下'c'键时，截取当前帧并允许用户选择ROI
            if cv2.waitKey(1) & 0xFF == ord('c'):
                # 转换到HSV
                hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
                roi = cv2.selectROI('Select color area', hsv)
                if roi != (0, 0, 0, 0):
                    # 计算ROI的平均颜色
                    roi_hsv = hsv[int(roi[1]):int(roi[1] + roi[3]), int(roi[0]):int(roi[0] + roi[2])]
                    color_mean = np.mean(roi_hsv, axis=(0, 1))
                    color_mean = np.uint8(color_mean)

                    color_lower = np.array(
                        [max(color_mean[0] - 30, 0), max(color_mean[1] - 60, 0), max(color_mean[2] - 55, 0)])
                    color_upper = np.array(
                        [min(color_mean[0] + 10, 179), 255, 255])
                    # print(f"Selected color: {color_mean}, Lower bound: {greenLower}, Upper bound: {greenUpper}")
                    color_selected = True
                    cv2.destroyWindow('Select color area')
                else:
                    print("No ROI selected. Exiting...")
                    break

        else:
            # 颜色区域已选择，进行颜色追踪
            # 将BGR图像转换为HSV
            hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)

            # 创建掩码
            mask = cv2.inRange(hsv, color_lower, color_upper)

            # 应用开运算和闭运算去除噪声
            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) > 300:
                    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('Original', frame)
            cv2.imshow('Mask', mask)

            # 按'q'键退出循环
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

    # 释放资源并关闭所有窗口
    cap.release()
    cv2.destroyAllWindows()

if __name__ == '__main__':
    main()