效果图

代码

import cv2

class VideoTracker:
    def __init__(self, video_path: str):
        self.video_path = video_path
        self.cap = cv2.VideoCapture(video_path)
        self.tracker = cv2.legacy.TrackerKCF_create()
        self.initBB = None
        self.tracker_initialized = False
        self.selecting = False
        self.frame = None
        self.paused = False
        self.tracking_started = False
        self.frame_skip = 2  # Adjusted for better tracking continuity

    def 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 = True
        elif 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 = False
            if self.initBB[2] > 0 and self.initBB[3] > 0:
                self.tracker.init(self.frame, self.initBB)
                self.tracker_initialized = True
                self.tracking_started = True

    def process_frame(self):
        """ Process frames for object tracking. """
        frame_counter = 0
        while self.cap.isOpened():
            ret, frame = self.cap.read()
            if not ret:
                break

            if 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-initializing
                            self.tracker_initialized = False
                            self.tracking_started = False
                    if not self.tracking_started:
                        # Optionally, re-select ROI here if desired
                        pass
                cv2.imshow('Video', self.frame)

            frame_counter += 1
            key = cv2.waitKey(1) & 0xFF
            if key == ord('q'):
                break
            elif key == ord(' '):  # Space key to pause/resume
                self.paused = not self.paused
            elif 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 = True
                    self.tracking_started = True

        self.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 ROI
        ret, self.frame = self.cap.read()
        if not ret:
            print("Error: Unable to read video file.")
            return

        self.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 ROI

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