|
|
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() |
