回传数据解析，兼容v5和v10

ultralytics/solutions/distance_calculation.py

@@ -0,0 +1,181 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+
+import math
+
+import cv2
+
+from ultralytics.utils.checks import check_imshow
+from ultralytics.utils.plotting import Annotator, colors
+
+
+class DistanceCalculation:
+    """A class to calculate distance between two objects in real-time video stream based on their tracks."""
+
+    def __init__(self):
+        """Initializes the distance calculation class with default values for Visual, Image, track and distance
+        parameters.
+        """
+
+        # Visual & im0 information
+        self.im0 = None
+        self.annotator = None
+        self.view_img = False
+        self.line_color = (255, 255, 0)
+        self.centroid_color = (255, 0, 255)
+
+        # Predict/track information
+        self.clss = None
+        self.names = None
+        self.boxes = None
+        self.line_thickness = 2
+        self.trk_ids = None
+
+        # Distance calculation information
+        self.centroids = []
+        self.pixel_per_meter = 10
+
+        # Mouse event
+        self.left_mouse_count = 0
+        self.selected_boxes = {}
+
+        # Check if environment support imshow
+        self.env_check = check_imshow(warn=True)
+
+    def set_args(
+        self,
+        names,
+        pixels_per_meter=10,
+        view_img=False,
+        line_thickness=2,
+        line_color=(255, 255, 0),
+        centroid_color=(255, 0, 255),
+    ):
+        """
+        Configures the distance calculation and display parameters.
+
+        Args:
+            names (dict): object detection classes names
+            pixels_per_meter (int): Number of pixels in meter
+            view_img (bool): Flag indicating frame display
+            line_thickness (int): Line thickness for bounding boxes.
+            line_color (RGB): color of centroids line
+            centroid_color (RGB): colors of bbox centroids
+        """
+        self.names = names
+        self.pixel_per_meter = pixels_per_meter
+        self.view_img = view_img
+        self.line_thickness = line_thickness
+        self.line_color = line_color
+        self.centroid_color = centroid_color
+
+    def mouse_event_for_distance(self, event, x, y, flags, param):
+        """
+        This function is designed to move region with mouse events in a real-time video stream.
+
+        Args:
+            event (int): The type of mouse event (e.g., cv2.EVENT_MOUSEMOVE, cv2.EVENT_LBUTTONDOWN, etc.).
+            x (int): The x-coordinate of the mouse pointer.
+            y (int): The y-coordinate of the mouse pointer.
+            flags (int): Any flags associated with the event (e.g., cv2.EVENT_FLAG_CTRLKEY,
+                cv2.EVENT_FLAG_SHIFTKEY, etc.).
+            param (dict): Additional parameters you may want to pass to the function.
+        """
+        global selected_boxes
+        global left_mouse_count
+        if event == cv2.EVENT_LBUTTONDOWN:
+            self.left_mouse_count += 1
+            if self.left_mouse_count <= 2:
+                for box, track_id in zip(self.boxes, self.trk_ids):
+                    if box[0] < x < box[2] and box[1] < y < box[3] and track_id not in self.selected_boxes:
+                        self.selected_boxes[track_id] = []
+                        self.selected_boxes[track_id] = box
+
+        if event == cv2.EVENT_RBUTTONDOWN:
+            self.selected_boxes = {}
+            self.left_mouse_count = 0
+
+    def extract_tracks(self, tracks):
+        """
+        Extracts results from the provided data.
+
+        Args:
+            tracks (list): List of tracks obtained from the object tracking process.
+        """
+        self.boxes = tracks[0].boxes.xyxy.cpu()
+        self.clss = tracks[0].boxes.cls.cpu().tolist()
+        self.trk_ids = tracks[0].boxes.id.int().cpu().tolist()
+
+    def calculate_centroid(self, box):
+        """
+        Calculate the centroid of bounding box.
+
+        Args:
+            box (list): Bounding box data
+        """
+        return int((box[0] + box[2]) // 2), int((box[1] + box[3]) // 2)
+
+    def calculate_distance(self, centroid1, centroid2):
+        """
+        Calculate distance between two centroids.
+
+        Args:
+            centroid1 (point): First bounding box data
+            centroid2 (point): Second bounding box data
+        """
+        pixel_distance = math.sqrt((centroid1[0] - centroid2[0]) ** 2 + (centroid1[1] - centroid2[1]) ** 2)
+        return pixel_distance / self.pixel_per_meter, (pixel_distance / self.pixel_per_meter) * 1000
+
+    def start_process(self, im0, tracks):
+        """
+        Calculate distance between two bounding boxes based on tracking data.
+
+        Args:
+            im0 (nd array): Image
+            tracks (list): List of tracks obtained from the object tracking process.
+        """
+        self.im0 = im0
+        if tracks[0].boxes.id is None:
+            if self.view_img:
+                self.display_frames()
+            return
+        self.extract_tracks(tracks)
+
+        self.annotator = Annotator(self.im0, line_width=2)
+
+        for box, cls, track_id in zip(self.boxes, self.clss, self.trk_ids):
+            self.annotator.box_label(box, color=colors(int(cls), True), label=self.names[int(cls)])
+
+            if len(self.selected_boxes) == 2:
+                for trk_id, _ in self.selected_boxes.items():
+                    if trk_id == track_id:
+                        self.selected_boxes[track_id] = box
+
+        if len(self.selected_boxes) == 2:
+            for trk_id, box in self.selected_boxes.items():
+                centroid = self.calculate_centroid(self.selected_boxes[trk_id])
+                self.centroids.append(centroid)
+
+            distance_m, distance_mm = self.calculate_distance(self.centroids[0], self.centroids[1])
+            self.annotator.plot_distance_and_line(
+                distance_m, distance_mm, self.centroids, self.line_color, self.centroid_color
+            )
+
+        self.centroids = []
+
+        if self.view_img and self.env_check:
+            self.display_frames()
+
+        return im0
+
+    def display_frames(self):
+        """Display frame."""
+        cv2.namedWindow("Ultralytics Distance Estimation")
+        cv2.setMouseCallback("Ultralytics Distance Estimation", self.mouse_event_for_distance)
+        cv2.imshow("Ultralytics Distance Estimation", self.im0)
+
+        if cv2.waitKey(1) & 0xFF == ord("q"):
+            return
+
+
+if __name__ == "__main__":
+    DistanceCalculation()