import cv2
import numpy as np

def detect_color_shape(
    image, 
    color_lower=np.array([0, 120, 70]), 
    color_upper=np.array([10, 255, 255]),
    shape_type="circle",
    min_area=100,
    max_area=10000,
    circularity_threshold=0.8,
    aspect_ratio_range=(0.9, 1.1)
):
    """
    检测指定颜色和形状的物体
    
    参数:
        image: 输入 RGB/BGR 图像
        color_lower: HSV 颜色下限
        color_upper: HSV 颜色上限
        shape_type: "circle" / "square" / "rectangle"
        min_area: 最小面积（像素）
        max_area: 最大面积（像素）
        circularity_threshold: 圆形度阈值（0~1，越大越圆）
        aspect_ratio_range: 宽高比范围（方形接近 1）
    
    返回:
        detections: 检测结果列表 [(x, y, w, h, contour), ...]
        mask: 颜色分割掩码（用于调试）
    """
    # 1. 颜色空间转换：BGR → HSV
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    
    # 2. 颜色阈值分割
    mask = cv2.inRange(hsv, color_lower, color_upper)
    
    # 3. 形态学处理去除噪声
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
    
    # 4. 查找轮廓
    contours, hierarchy = cv2.findContours(
        mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
    )
    
    detections = []
    
    for contour in contours:
        # 计算轮廓面积，过滤太小/太大的
        area = cv2.contourArea(contour)
        if area < min_area or area > max_area:
            continue
        
        # 计算轮廓周长
        perimeter = cv2.arcLength(contour, True)
        
        # 多边形逼近
        approx = cv2.approxPolyDP(contour, 0.04 * perimeter, True)
        
        # 获取外接矩形
        x, y, w, h = cv2.boundingRect(contour)
        
        # 形状判断
        matched = False
        
        if shape_type == "circle":
            # 圆形度 = 4π * 面积 / (周长^2)
            # 完美圆形 = 1，越不规则值越小
            circularity = 4 * np.pi * area / (perimeter * perimeter)
            if circularity >= circularity_threshold:
                matched = True
        
        elif shape_type == "square":
            # 宽高比接近 1，且顶点数约为 4
            aspect_ratio = float(w) / h
            if (len(approx) == 4 and 
                aspect_ratio >= aspect_ratio_range[0] and 
                aspect_ratio <= aspect_ratio_range[1]):
                matched = True
        
        elif shape_type == "rectangle":
            # 顶点数约为 4 即可
            if len(approx) == 4:
                matched = True
        
        elif shape_type == "triangle":
            if len(approx) == 3:
                matched = True
        
        if matched:
            center_x = x + w // 2
            center_y = y + h // 2
            detections.append((center_x, center_y, w, h, contour))
    
    return detections, mask

# 检测红色圆形物体
# HSV 红色范围（两种分段，因为红色在 H 通道首尾）
lower_red1 = np.array([0, 120, 70])
upper_red1 = np.array([10, 255, 255])
lower_red2 = np.array([170, 120, 70])
upper_red2 = np.array([180, 255, 255])

# 读取图像
image = cv2.imread("test_image.jpg")

# 第一次检测
detections1, mask1 = detect_color_shape(
    image, lower_red1, upper_red1, 
    shape_type="circle", min_area=500
)

# 第二次检测（另一红色区间）
detections2, mask2 = detect_color_shape(
    image, lower_red2, upper_red2, 
    shape_type="circle", min_area=500
)

# 合并结果
all_detections = detections1 + detections2
combined_mask = mask1 | mask2

# 在原图上绘制结果
for (x, y, w, h, contour) in all_detections:
    cv2.drawContours(image, [contour], -1, (0, 255, 0), 2)
    cv2.circle(image, (x, y), 3, (0, 0, 255), -1)
    cv2.putText(image, "Red Circle", (x-20, y-h-10),
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

cv2.imwrite("result.jpg", image)
print(f"检测到 {len(all_detections)} 个红色圆形物体")

def create_hsv_trackbars(window_name="HSV Tuner"):
    cv2.namedWindow(window_name)
    cv2.createTrackbar("H_min", window_name, 0, 180, lambda x: None)
    cv2.createTrackbar("H_max", window_name, 180, 180, lambda x: None)
    cv2.createTrackbar("S_min", window_name, 0, 255, lambda x: None)
    cv2.createTrackbar("S_max", window_name, 255, 255, lambda x: None)
    cv2.createTrackbar("V_min", window_name, 0, 255, lambda x: None)
    cv2.createTrackbar("V_max", window_name, 255, 255, lambda x: None)

def get_hsv_from_trackbars(window_name="HSV Tuner"):
    h_min = cv2.getTrackbarPos("H_min", window_name)
    h_max = cv2.getTrackbarPos("H_max", window_name)
    s_min = cv2.getTrackbarPos("S_min", window_name)
    s_max = cv2.getTrackbarPos("S_max", window_name)
    v_min = cv2.getTrackbarPos("V_min", window_name)
    v_max = cv2.getTrackbarPos("V_max", window_name)
    return np.array([h_min, s_min, v_min]), np.array([h_max, s_max, v_max])

def detect_circles_hough(
    image, 
    dp=1.2, 
    min_dist=50, 
    param1=50, 
    param2=30, 
    min_radius=10, 
    max_radius=100
):
    """
    使用 Hough 变换检测圆形
    
    参数说明:
        dp: 累加器分辨率反比，1 = 原图分辨率
        min_dist: 两个圆心之间的最小距离（避免重复检测）
        param1: Canny 边缘检测高阈值
        param2: 累加器阈值，越大检测越少但越准
        min_radius/max_radius: 半径范围
    
    返回:
        circles: [(x, y, r), ...]
    """
    # 转灰度
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    
    # 高斯模糊降噪
    gray = cv2.GaussianBlur(gray, (5, 5), 0)
    
    # Hough 变换检测圆
    circles = cv2.HoughCircles(
        gray, 
        cv2.HOUGH_GRADIENT, 
        dp=dp, 
        minDist=min_dist,
        param1=param1, 
        param2=param2, 
        minRadius=min_radius, 
        maxRadius=max_radius
    )
    
    result = []
    if circles is not None:
        circles = np.uint16(np.around(circles))
        for i in circles[0, :]:
            x, y, r = i[0], i[1], i[2]
            result.append((x, y, r))
    
    return result

# 使用示例
image = cv2.imread("curling.jpg")
circles = detect_circles_hough(
    image, 
    min_radius=80, 
    max_radius=120  # 冰壶直径约 100 像素
)

# 绘制结果
for (x, y, r) in circles:
    cv2.circle(image, (x, y), r, (0, 255, 0), 2)  # 圆周
    cv2.circle(image, (x, y), 3, (0, 0, 255), -1)  # 圆心

print(f"检测到 {len(circles)} 个冰壶")

def detect_blobs_by_color_and_shape(
    image,
    color_lower,
    color_upper,
    min_area=100,
    max_area=10000,
    min_circularity=0.5,
    max_circularity=1.0,
    min_inertia_ratio=0.5,
    min_convexity=0.5
):
    """
    基于 Blob 分析检测指定颜色和形状的物体
    
    参数:
        min_inertia_ratio: 惯性比，圆 ≈ 1，长条 ≈ 0
        min_convexity: 凸包面积比，越凸越接近 1
    """
    # 颜色分割得到掩码
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    mask = cv2.inRange(hsv, color_lower, color_upper)
    
    # 设置 Blob 检测器参数
    params = cv2.SimpleBlobDetector_Params()
    
    # 按面积过滤
    params.filterByArea = True
    params.minArea = min_area
    params.maxArea = max_area
    
    # 按圆形度过滤
    params.filterByCircularity = True
    params.minCircularity = min_circularity
    params.maxCircularity = max_circularity
    
    # 按惯性比过滤（判断是否圆）
    params.filterByInertia = True
    params.minInertiaRatio = min_inertia_ratio
    
    # 按凸性过滤
    params.filterByConvexity = True
    params.minConvexity = min_convexity
    
    # 创建检测器并检测
    detector = cv2.SimpleBlobDetector_create(params)
    keypoints = detector.detect(mask)
    
    # 提取结果
    result = []
    for kp in keypoints:
        x = int(kp.pt[0])
        y = int(kp.pt[1])
        size = int(kp.size)
        r = int(size / 2)
        result.append((x, y, r))
    
    return result, mask, keypoints

pip install ultralytics

from ultralytics import YOLO

# 加载预训练模型（或你自己训练的模型）
model = YOLO("yolov8n.pt")  # 检测自定义类需要重新训练

# 假设我们训练了一个检测 "red_circle", "blue_square" 的模型
# model = YOLO("best.pt")  # 你的训练结果

def detect_objects_yolo(image, conf_threshold=0.5):
    """YOLOv8 目标检测"""
    results = model(image, conf=conf_threshold)[0]
    
    detections = []
    for box in results.boxes:
        x1, y1, x2, y2 = box.xyxy[0].tolist()
        conf = box.conf[0].item()
        class_id = int(box.cls[0].item())
        class_name = model.names[class_id]
        
        x = int((x1 + x2) / 2)
        y = int((y1 + y2) / 2)
        w = int(x2 - x1)
        h = int(y2 - y1)
        
        detections.append({
            "bbox": (x, y, w, h),
            "class_name": class_name,
            "confidence": conf
        })
    
    return detections

# 使用示例
image = cv2.imread("test.jpg")
detections = detect_objects_yolo(image, conf_threshold=0.5)

for det in detections:
    x, y, w, h = det["bbox"]
    name = det["class_name"]
    conf = det["confidence"]
    
    x1 = x - w // 2
    y1 = y - h // 2
    x2 = x + w // 2
    y2 = y + h // 2
    
    cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
    cv2.putText(image, f"{name} {conf:.2f}", (x1, y1-10),
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

print(f"检测到 {len(detections)} 个目标")

# 冰壶检测完整示例
def detect_curling_stones(image, min_radius=60, max_radius=100):
    # 冰壶是深色，冰面是浅色
    # 转灰度 + 阈值分割
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    _, mask = cv2.threshold(gray, 100, 255, cv2.THRESH_BINARY_INV)
    
    # 形态学处理
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
    
    # Hough 检测圆
    circles = cv2.HoughCircles(
        mask, cv2.HOUGH_GRADIENT, 
        dp=1.2, min_dist=80,
        param1=50, param2=25,
        minRadius=min_radius, maxRadius=max_radius
    )
    
    result = []
    if circles is not None:
        circles = np.uint16(np.around(circles))
        for i in circles[0, :]:
            result.append((i[0], i[1], i[2]))
    
    return result, mask