import cv2
import numpy as np

cap = cv2.VideoCapture(0)

cap.set(3, 320)
cap.set(4, 240)
cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)

# ===== 核心改动：用背景减法器替代帧差法 =====
# MOG2: 能适应背景变化，适合移动摄像头场景
# history: 保留多少帧来建立背景模型（越大适应越慢，但越稳定）
# varThreshold: 判断前景的灵敏度（越小越敏感，但噪声越多）
fgbg = cv2.createBackgroundSubtractorMOG2(history=100, varThreshold=25, detectShadows=False)

# 可选：用KNN背景减法器（效果类似，速度稍快）
# fgbg = cv2.createBackgroundSubtractorKNN(history=100, dist2Threshold=400, detectShadows=False)

# 形态学核（用于去噪）
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))

# 用于平滑追踪的卡尔曼滤波器（可选，减少抖动）
kalman = cv2.KalmanFilter(4, 2)
kalman.measurementMatrix = np.array([[1,0,0,0], [0,1,0,0]], np.float32)
kalman.transitionMatrix = np.array([[1,0,1,0], [0,1,0,1], [0,0,1,0], [0,0,0,1]], np.float32)
kalman.processNoiseCov = np.eye(4, dtype=np.float32) * 0.03
kalman.measurementNoiseCov = np.eye(2, dtype=np.float32) * 0.5
kalman.statePost = np.zeros((4,1), np.float32)

while True:
    ret, frame = cap.read()
    if not ret:
        break
    
    # 1. 背景减法：得到前景掩码（白色=运动物体）
    fgmask = fgbg.apply(frame)
    
    # 2. 形态学处理：去噪 + 填充空洞
    fgmask = cv2.morphologyEx(fgmask, cv2.MORPH_OPEN, kernel)   # 先开运算去噪点
    fgmask = cv2.morphologyEx(fgmask, cv2.MORPH_CLOSE, kernel)  # 再闭运算填充空洞
    
    # 3. 找轮廓
    contours, _ = cv2.findContours(fgmask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    # ===== 核心：选最大目标 =====
    target = None
    max_area = 0
    
    for cnt in contours:
        area = cv2.contourArea(cnt)
        
        if area < 300:  # 过滤噪声
            continue
        
        if area > max_area:
            max_area = area
            target = cnt
    
    # ===== 处理目标 =====
    if target is not None:
        x, y, w, h = cv2.boundingRect(target)
        cx = x + w // 2
        cy = y + h // 2
        
        # 卡尔曼滤波预测/更新（可选，让追踪更平滑）
        measurement = np.array([[np.float32(cx)], [np.float32(cy)]])
        kalman.correct(measurement)
        prediction = kalman.predict()
        pred_x, pred_y = int(prediction[0]), int(prediction[1])
        
        # 画原始检测框
        cv2.rectangle(frame, (x, y), (x+w, y+h), (0, 255, 0), 2)
        cv2.circle(frame, (cx, cy), 5, (0, 0, 255), -1)
        # 画卡尔曼平滑后的位置（黄色）
        cv2.circle(frame, (pred_x, pred_y), 5, (0, 255, 255), -1)
        
        print(f"目标位置: ({cx}, {cy}), 面积: {max_area}")
    else:
        # 没有检测到目标时，只预测不更新
        prediction = kalman.predict()
        # 可选：根据预测位置保持追踪（即使被短暂遮挡）
    
    # 可选：显示前景掩码（调试用）
    cv2.imshow("foreground mask", fgmask)
    cv2.imshow("tracking", frame)
    
    if cv2.waitKey(1) & 0xFF == 27:
        break

cap.release()
cv2.destroyAllWindows()