# -*- coding:utf-8 -*-
# @Author len
# @Create 2023/11/22 16:30

import cv2
import numpy as np

def get_color_threshold(image_path, color):
    # 加载图像
    image = cv2.imread(image_path)

    # 将图像转换为LAB颜色空间
    lab_image = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)

    # 根据不同的颜色选择相应的阈值
    if color == "red":
        lower_threshold = np.array([0, 127, 127])
        upper_threshold = np.array([30, 255, 255])
    elif color == "green":
        lower_threshold = np.array([30, -8, -32])
        upper_threshold = np.array([100, 64, 32])
    elif color == "blue":
        lower_threshold = np.array([0, 0, -128])
        upper_threshold = np.array([30, 64, -20])
    else:
        raise ValueError("Invalid color")

    # 应用阈值，获取掩码
    mask = cv2.inRange(lab_image, lower_threshold, upper_threshold)

    # 返回阈值掩码
    return mask

# 示例用法
image_path = "../first_frame1.jpg"  # 替换为您的图像路径
color = "red"  # 替换为您想要获取阈值的颜色
threshold_mask = get_color_threshold(image_path, color)

# 显示阈值图像
cv2.imshow("Threshold", threshold_mask)
cv2.waitKey(0)
cv2.destroyAllWindows()