Bỏ để qua phần nội dung

Tài liệu tham khảo cho ultralytics/solutions/distance_calculation.py

Ghi

Tệp này có sẵn tại https://github.com/ultralytics/ultralytics/blob/main/ultralytics/giải pháp/distance_calculation.py. Nếu bạn phát hiện ra một vấn đề, vui lòng giúp khắc phục nó bằng cách đóng góp Yêu cầu 🛠️ kéo. Cảm ơn bạn 🙏 !



ultralytics.solutions.distance_calculation.DistanceCalculation

Một lớp để tính toán khoảng cách giữa hai đối tượng trong luồng video thời gian thực dựa trên các bản nhạc của chúng.

Mã nguồn trong ultralytics/solutions/distance_calculation.py
class DistanceCalculation:
    """A class to calculate distance between two objects in a real-time video stream based on their tracks."""

    def __init__(
        self,
        names,
        pixels_per_meter=10,
        view_img=False,
        line_thickness=2,
        line_color=(255, 255, 0),
        centroid_color=(255, 0, 255),
    ):
        """
        Initializes the DistanceCalculation class with the given parameters.

        Args:
            names (dict): Dictionary mapping class indices to class names.
            pixels_per_meter (int, optional): Conversion factor from pixels to meters. Defaults to 10.
            view_img (bool, optional): Flag to indicate if the video stream should be displayed. Defaults to False.
            line_thickness (int, optional): Thickness of the lines drawn on the image. Defaults to 2.
            line_color (tuple, optional): Color of the lines drawn on the image (BGR format). Defaults to (255, 255, 0).
            centroid_color (tuple, optional): Color of the centroids drawn (BGR format). Defaults to (255, 0, 255).
        """
        # Visual & image information
        self.im0 = None
        self.annotator = None
        self.view_img = view_img
        self.line_color = line_color
        self.centroid_color = centroid_color

        # Prediction & tracking information
        self.clss = None
        self.names = names
        self.boxes = None
        self.line_thickness = line_thickness
        self.trk_ids = None

        # Distance calculation information
        self.centroids = []
        self.pixel_per_meter = pixels_per_meter

        # Mouse event information
        self.left_mouse_count = 0
        self.selected_boxes = {}

        # Check if environment supports imshow
        self.env_check = check_imshow(warn=True)

    def mouse_event_for_distance(self, event, x, y, flags, param):
        """
        Handles mouse events to select regions in a real-time video stream.

        Args:
            event (int): Type of mouse event (e.g., cv2.EVENT_MOUSEMOVE, cv2.EVENT_LBUTTONDOWN, etc.).
            x (int): X-coordinate of the mouse pointer.
            y (int): Y-coordinate of the mouse pointer.
            flags (int): Flags associated with the event (e.g., cv2.EVENT_FLAG_CTRLKEY, cv2.EVENT_FLAG_SHIFTKEY, etc.).
            param (dict): Additional parameters passed to the function.
        """
        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] = box

        elif event == cv2.EVENT_RBUTTONDOWN:
            self.selected_boxes = {}
            self.left_mouse_count = 0

    def extract_tracks(self, tracks):
        """
        Extracts tracking 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()

    @staticmethod
    def calculate_centroid(box):
        """
        Calculates the centroid of a bounding box.

        Args:
            box (list): Bounding box coordinates [x1, y1, x2, y2].

        Returns:
            (tuple): Centroid coordinates (x, y).
        """
        return int((box[0] + box[2]) // 2), int((box[1] + box[3]) // 2)

    def calculate_distance(self, centroid1, centroid2):
        """
        Calculates the distance between two centroids.

        Args:
            centroid1 (tuple): Coordinates of the first centroid (x, y).
            centroid2 (tuple): Coordinates of the second centroid (x, y).

        Returns:
            (tuple): Distance in meters and millimeters.
        """
        pixel_distance = math.sqrt((centroid1[0] - centroid2[0]) ** 2 + (centroid1[1] - centroid2[1]) ** 2)
        distance_m = pixel_distance / self.pixel_per_meter
        distance_mm = distance_m * 1000
        return distance_m, distance_mm

    def start_process(self, im0, tracks):
        """
        Processes the video frame and calculates the distance between two bounding boxes.

        Args:
            im0 (ndarray): The image frame.
            tracks (list): List of tracks obtained from the object tracking process.

        Returns:
            (ndarray): The processed image frame.
        """
        self.im0 = im0
        if tracks[0].boxes.id is None:
            if self.view_img:
                self.display_frames()
            return im0

        self.extract_tracks(tracks)
        self.annotator = Annotator(self.im0, line_width=self.line_thickness)

        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.keys():
                    if trk_id == track_id:
                        self.selected_boxes[track_id] = box

        if len(self.selected_boxes) == 2:
            self.centroids = [self.calculate_centroid(self.selected_boxes[trk_id]) for trk_id in self.selected_boxes]

            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):
        """Displays the current frame with annotations."""
        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

__init__(names, pixels_per_meter=10, view_img=False, line_thickness=2, line_color=(255, 255, 0), centroid_color=(255, 0, 255))

Khởi tạo lớp DistanceCalculation với các tham số đã cho.

Thông số:

Tên Kiểu Sự miêu tả Mặc định
names dict

Từ điển ánh xạ các chỉ mục lớp đến tên lớp.

bắt buộc
pixels_per_meter int

Hệ số chuyển đổi từ pixel sang mét. Mặc định là 10.

10
view_img bool

Gắn cờ để cho biết luồng video có hiển thị hay không. Mặc định là False.

False
line_thickness int

Độ dày của các đường vẽ trên hình ảnh. Mặc định là 2.

2
line_color tuple

Màu sắc của các đường vẽ trên hình ảnh (định dạng BGR). Mặc định là (255, 255, 0).

(255, 255, 0)
centroid_color tuple

Màu sắc của centroid được vẽ (định dạng BGR). Mặc định là (255, 0, 255).

(255, 0, 255)
Mã nguồn trong ultralytics/solutions/distance_calculation.py
def __init__(
    self,
    names,
    pixels_per_meter=10,
    view_img=False,
    line_thickness=2,
    line_color=(255, 255, 0),
    centroid_color=(255, 0, 255),
):
    """
    Initializes the DistanceCalculation class with the given parameters.

    Args:
        names (dict): Dictionary mapping class indices to class names.
        pixels_per_meter (int, optional): Conversion factor from pixels to meters. Defaults to 10.
        view_img (bool, optional): Flag to indicate if the video stream should be displayed. Defaults to False.
        line_thickness (int, optional): Thickness of the lines drawn on the image. Defaults to 2.
        line_color (tuple, optional): Color of the lines drawn on the image (BGR format). Defaults to (255, 255, 0).
        centroid_color (tuple, optional): Color of the centroids drawn (BGR format). Defaults to (255, 0, 255).
    """
    # Visual & image information
    self.im0 = None
    self.annotator = None
    self.view_img = view_img
    self.line_color = line_color
    self.centroid_color = centroid_color

    # Prediction & tracking information
    self.clss = None
    self.names = names
    self.boxes = None
    self.line_thickness = line_thickness
    self.trk_ids = None

    # Distance calculation information
    self.centroids = []
    self.pixel_per_meter = pixels_per_meter

    # Mouse event information
    self.left_mouse_count = 0
    self.selected_boxes = {}

    # Check if environment supports imshow
    self.env_check = check_imshow(warn=True)

calculate_centroid(box) staticmethod

Tính tâm của một hộp giới hạn.

Thông số:

Tên Kiểu Sự miêu tả Mặc định
box list

Tọa độ hộp giới hạn [x1, y1, x2, y2].

bắt buộc

Trở lại:

Kiểu Sự miêu tả
tuple

Tọa độ trung tâm (x, y).

Mã nguồn trong ultralytics/solutions/distance_calculation.py
@staticmethod
def calculate_centroid(box):
    """
    Calculates the centroid of a bounding box.

    Args:
        box (list): Bounding box coordinates [x1, y1, x2, y2].

    Returns:
        (tuple): Centroid coordinates (x, y).
    """
    return int((box[0] + box[2]) // 2), int((box[1] + box[3]) // 2)

calculate_distance(centroid1, centroid2)

Tính khoảng cách giữa hai centroid.

Thông số:

Tên Kiểu Sự miêu tả Mặc định
centroid1 tuple

Tọa độ của trung tâm thứ nhất (x, y).

bắt buộc
centroid2 tuple

Tọa độ của trung tâm thứ hai (x, y).

bắt buộc

Trở lại:

Kiểu Sự miêu tả
tuple

Khoảng cách tính bằng mét và milimét.

Mã nguồn trong ultralytics/solutions/distance_calculation.py
def calculate_distance(self, centroid1, centroid2):
    """
    Calculates the distance between two centroids.

    Args:
        centroid1 (tuple): Coordinates of the first centroid (x, y).
        centroid2 (tuple): Coordinates of the second centroid (x, y).

    Returns:
        (tuple): Distance in meters and millimeters.
    """
    pixel_distance = math.sqrt((centroid1[0] - centroid2[0]) ** 2 + (centroid1[1] - centroid2[1]) ** 2)
    distance_m = pixel_distance / self.pixel_per_meter
    distance_mm = distance_m * 1000
    return distance_m, distance_mm

display_frames()

Hiển thị khung hiện tại với chú thích.

Mã nguồn trong ultralytics/solutions/distance_calculation.py
def display_frames(self):
    """Displays the current frame with annotations."""
    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

extract_tracks(tracks)

Trích xuất kết quả theo dõi từ dữ liệu được cung cấp.

Thông số:

Tên Kiểu Sự miêu tả Mặc định
tracks list

Danh sách các bản nhạc thu được từ quá trình theo dõi đối tượng.

bắt buộc
Mã nguồn trong ultralytics/solutions/distance_calculation.py
def extract_tracks(self, tracks):
    """
    Extracts tracking 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()

mouse_event_for_distance(event, x, y, flags, param)

Xử lý các sự kiện chuột để chọn các vùng trong luồng video thời gian thực.

Thông số:

Tên Kiểu Sự miêu tả Mặc định
event int

Loại sự kiện chuột (ví dụ: cv2. EVENT_MOUSEMOVE, CV2. EVENT_LBUTTONDOWN, v.v.).

bắt buộc
x int

Tọa độ X của con trỏ chuột.

bắt buộc
y int

Tọa độ Y của con trỏ chuột.

bắt buộc
flags int

Cờ liên quan đến sự kiện (ví dụ: cv2. EVENT_FLAG_CTRLKEY, CV2. EVENT_FLAG_SHIFTKEY, v.v.).

bắt buộc
param dict

Các tham số bổ sung được chuyển đến hàm.

bắt buộc
Mã nguồn trong ultralytics/solutions/distance_calculation.py
def mouse_event_for_distance(self, event, x, y, flags, param):
    """
    Handles mouse events to select regions in a real-time video stream.

    Args:
        event (int): Type of mouse event (e.g., cv2.EVENT_MOUSEMOVE, cv2.EVENT_LBUTTONDOWN, etc.).
        x (int): X-coordinate of the mouse pointer.
        y (int): Y-coordinate of the mouse pointer.
        flags (int): Flags associated with the event (e.g., cv2.EVENT_FLAG_CTRLKEY, cv2.EVENT_FLAG_SHIFTKEY, etc.).
        param (dict): Additional parameters passed to the function.
    """
    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] = box

    elif event == cv2.EVENT_RBUTTONDOWN:
        self.selected_boxes = {}
        self.left_mouse_count = 0

start_process(im0, tracks)

Xử lý khung hình video và tính toán khoảng cách giữa hai hộp giới hạn.

Thông số:

Tên Kiểu Sự miêu tả Mặc định
im0 ndarray

Khung ảnh.

bắt buộc
tracks list

Danh sách các bản nhạc thu được từ quá trình theo dõi đối tượng.

bắt buộc

Trở lại:

Kiểu Sự miêu tả
ndarray

Khung hình ảnh được xử lý.

Mã nguồn trong ultralytics/solutions/distance_calculation.py
def start_process(self, im0, tracks):
    """
    Processes the video frame and calculates the distance between two bounding boxes.

    Args:
        im0 (ndarray): The image frame.
        tracks (list): List of tracks obtained from the object tracking process.

    Returns:
        (ndarray): The processed image frame.
    """
    self.im0 = im0
    if tracks[0].boxes.id is None:
        if self.view_img:
            self.display_frames()
        return im0

    self.extract_tracks(tracks)
    self.annotator = Annotator(self.im0, line_width=self.line_thickness)

    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.keys():
                if trk_id == track_id:
                    self.selected_boxes[track_id] = box

    if len(self.selected_boxes) == 2:
        self.centroids = [self.calculate_centroid(self.selected_boxes[trk_id]) for trk_id in self.selected_boxes]

        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





Created 2024-01-05, Updated 2024-06-02
Authors: glenn-jocher (2), Burhan-Q (1), AyushExel (1), RizwanMunawar (1)