import cv2
import numpy as np

def main():
    # 初始化摄像头
    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_range = 10
                    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)
            print(color_lower, color_upper)

            # 显示原图和结果图像
            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()