效果图
代码
import cv2class VideoTracker:def __init__(self, video_path: str):self.video_path = video_pathself.cap = cv2.VideoCapture(video_path)self.tracker = cv2.legacy.TrackerKCF_create()self.initBB = Noneself.tracker_initialized = Falseself.selecting = Falseself.frame = Noneself.paused = Falseself.tracking_started = Falseself.frame_skip = 2 # Adjusted for better tracking continuitydef select_roi(self, event, x, y, flags, param):""" Handle mouse events for ROI selection. """if event == cv2.EVENT_LBUTTONDOWN:self.initBB = (x, y, 0, 0)self.selecting = Trueelif event == cv2.EVENT_MOUSEMOVE and self.selecting:self.initBB = (self.initBB[0], self.initBB[1], x - self.initBB[0], y - self.initBB[1])self.frame_copy = self.frame.copy()cv2.rectangle(self.frame_copy, (self.initBB[0], self.initBB[1]), (x, y), (255, 0, 0), 2)cv2.imshow('Video', self.frame_copy)elif event == cv2.EVENT_LBUTTONUP:self.initBB = (self.initBB[0], self.initBB[1], x - self.initBB[0], y - self.initBB[1])self.selecting = Falseif self.initBB[2] > 0 and self.initBB[3] > 0:self.tracker.init(self.frame, self.initBB)self.tracker_initialized = Trueself.tracking_started = Truedef process_frame(self):""" Process frames for object tracking. """frame_counter = 0while self.cap.isOpened():ret, frame = self.cap.read()if not ret:breakif frame_counter % (self.frame_skip + 1) == 0:self.frame = frame.copy()if self.tracking_started:if self.tracker_initialized:success, bbox = self.tracker.update(self.frame)if success:(x, y, w, h) = [int(v) for v in bbox]center_x, center_y = int(x + w / 2), int(y + h / 2)cv2.rectangle(self.frame, (x, y), (x + w, y + h), (0, 255, 0), 2)cv2.circle(self.frame, (center_x, center_y), 5, (0, 0, 255), -1)cv2.putText(self.frame, f"Center: ({center_x}, {center_y})", (center_x + 10, center_y - 10),cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 2)else:# Tracker failed, try re-initializingself.tracker_initialized = Falseself.tracking_started = Falseif not self.tracking_started:# Optionally, re-select ROI here if desiredpasscv2.imshow('Video', self.frame)frame_counter += 1key = cv2.waitKey(1) & 0xFFif key == ord('q'):breakelif key == ord(' '): # Space key to pause/resumeself.paused = not self.pausedelif key == ord('s') and not self.selecting and not self.tracking_started:if self.initBB:self.tracker.init(self.frame, self.initBB)self.tracker_initialized = Trueself.tracking_started = Trueself.cap.release()cv2.destroyAllWindows()def start(self):""" Start the video tracking process. """cv2.namedWindow('Video')cv2.setMouseCallback('Video', self.select_roi)# Read the first frame and let the user select the ROIret, self.frame = self.cap.read()if not ret:print("Error: Unable to read video file.")returnself.frame_copy = self.frame.copy()cv2.imshow('Video', self.frame)print("Select ROI and press 's' to start tracking.")cv2.waitKey(0) # Wait for the user to select ROIself.process_frame()if __name__ == '__main__':video_path = 'ego-planner户外飞行测试.mp4' # 修改为你的实际视频路径if not video_path:print("Error: No video path provided.")else:tracker = VideoTracker(video_path)tracker.start()