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

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

Ghi

Tệp này có sẵn tại https://github.com/ultralytics/ultralytics/blob/main/ultralytics/giải pháp/queue_management.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.queue_management.QueueManager

Một lớp để quản lý hàng đợi trong luồng video thời gian thực dựa trên các bản nhạc đối tượng.

Mã nguồn trong ultralytics/solutions/queue_management.py
class QueueManager:
    """A class to manage the queue in a real-time video stream based on object tracks."""

    def __init__(
        self,
        classes_names,
        reg_pts=None,
        line_thickness=2,
        track_thickness=2,
        view_img=False,
        region_color=(255, 0, 255),
        view_queue_counts=True,
        draw_tracks=False,
        count_txt_color=(255, 255, 255),
        track_color=None,
        region_thickness=5,
        fontsize=0.7,
    ):
        """
        Initializes the QueueManager with specified parameters for tracking and counting objects.

        Args:
            classes_names (dict): A dictionary mapping class IDs to class names.
            reg_pts (list of tuples, optional): Points defining the counting region polygon. Defaults to a predefined
                rectangle.
            line_thickness (int, optional): Thickness of the annotation lines. Defaults to 2.
            track_thickness (int, optional): Thickness of the track lines. Defaults to 2.
            view_img (bool, optional): Whether to display the image frames. Defaults to False.
            region_color (tuple, optional): Color of the counting region lines (BGR). Defaults to (255, 0, 255).
            view_queue_counts (bool, optional): Whether to display the queue counts. Defaults to True.
            draw_tracks (bool, optional): Whether to draw tracks of the objects. Defaults to False.
            count_txt_color (tuple, optional): Color of the count text (BGR). Defaults to (255, 255, 255).
            track_color (tuple, optional): Color of the tracks. If None, different colors will be used for different
                tracks. Defaults to None.
            region_thickness (int, optional): Thickness of the counting region lines. Defaults to 5.
            fontsize (float, optional): Font size for the text annotations. Defaults to 0.7.
        """

        # Mouse events state
        self.is_drawing = False
        self.selected_point = None

        # Region & Line Information
        self.reg_pts = reg_pts if reg_pts is not None else [(20, 60), (20, 680), (1120, 680), (1120, 60)]
        self.counting_region = (
            Polygon(self.reg_pts) if len(self.reg_pts) >= 3 else Polygon([(20, 60), (20, 680), (1120, 680), (1120, 60)])
        )
        self.region_color = region_color
        self.region_thickness = region_thickness

        # Image and annotation Information
        self.im0 = None
        self.tf = line_thickness
        self.view_img = view_img
        self.view_queue_counts = view_queue_counts
        self.fontsize = fontsize

        self.names = classes_names  # Class names
        self.annotator = None  # Annotator
        self.window_name = "Ultralytics YOLOv8 Queue Manager"

        # Object counting Information
        self.counts = 0
        self.count_txt_color = count_txt_color

        # Tracks info
        self.track_history = defaultdict(list)
        self.track_thickness = track_thickness
        self.draw_tracks = draw_tracks
        self.track_color = track_color

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

    def extract_and_process_tracks(self, tracks):
        """Extracts and processes tracks for queue management in a video stream."""

        # Initialize annotator and draw the queue region
        self.annotator = Annotator(self.im0, self.tf, self.names)

        if tracks[0].boxes.id is not None:
            boxes = tracks[0].boxes.xyxy.cpu()
            clss = tracks[0].boxes.cls.cpu().tolist()
            track_ids = tracks[0].boxes.id.int().cpu().tolist()

            # Extract tracks
            for box, track_id, cls in zip(boxes, track_ids, clss):
                # Draw bounding box
                self.annotator.box_label(box, label=f"{self.names[cls]}#{track_id}", color=colors(int(track_id), True))

                # Update track history
                track_line = self.track_history[track_id]
                track_line.append((float((box[0] + box[2]) / 2), float((box[1] + box[3]) / 2)))
                if len(track_line) > 30:
                    track_line.pop(0)

                # Draw track trails if enabled
                if self.draw_tracks:
                    self.annotator.draw_centroid_and_tracks(
                        track_line,
                        color=self.track_color or colors(int(track_id), True),
                        track_thickness=self.track_thickness,
                    )

                prev_position = self.track_history[track_id][-2] if len(self.track_history[track_id]) > 1 else None

                # Check if the object is inside the counting region
                if len(self.reg_pts) >= 3:
                    is_inside = self.counting_region.contains(Point(track_line[-1]))
                    if prev_position is not None and is_inside:
                        self.counts += 1

        # Display queue counts
        label = f"Queue Counts : {str(self.counts)}"
        if label is not None:
            self.annotator.queue_counts_display(
                label,
                points=self.reg_pts,
                region_color=self.region_color,
                txt_color=self.count_txt_color,
            )

        self.counts = 0  # Reset counts after displaying
        self.display_frames()

    def display_frames(self):
        """Displays the current frame with annotations."""
        if self.env_check:
            self.annotator.draw_region(reg_pts=self.reg_pts, thickness=self.region_thickness, color=self.region_color)
            cv2.namedWindow(self.window_name)
            cv2.imshow(self.window_name, self.im0)
            # Close window on 'q' key press
            if cv2.waitKey(1) & 0xFF == ord("q"):
                return

    def process_queue(self, im0, tracks):
        """
        Main function to start the queue management process.

        Args:
            im0 (ndarray): Current frame from the video stream.
            tracks (list): List of tracks obtained from the object tracking process.
        """
        self.im0 = im0  # Store the current frame
        self.extract_and_process_tracks(tracks)  # Extract and process tracks

        if self.view_img:
            self.display_frames()  # Display the frame if enabled
        return self.im0

__init__(classes_names, reg_pts=None, line_thickness=2, track_thickness=2, view_img=False, region_color=(255, 0, 255), view_queue_counts=True, draw_tracks=False, count_txt_color=(255, 255, 255), track_color=None, region_thickness=5, fontsize=0.7)

Khởi tạo QueueManager với các tham số được chỉ định để theo dõi và đếm đối tượng.

Thông số:

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

Từ điển ánh xạ ID lớp đến tên lớp.

bắt buộc
reg_pts list of tuples

Các điểm xác định đa giác vùng đếm. Mặc định là một xác định trước hình chữ nhật.

None
line_thickness int

Độ dày của các dòng chú thích. Mặc định là 2.

2
track_thickness int

Độ dày của các đường ray. Mặc định là 2.

2
view_img bool

Có hiển thị khung hình ảnh hay không. Mặc định là False.

False
region_color tuple

Màu sắc của các đường vùng đếm (BGR). Mặc định là (255, 0, 255).

(255, 0, 255)
view_queue_counts bool

Có hiển thị số lượng hàng đợi hay không. Mặc định là True.

True
draw_tracks bool

Có nên vẽ dấu vết của các đối tượng hay không. Mặc định là False.

False
count_txt_color tuple

Màu của văn bản đếm (BGR). Mặc định là (255, 255, 255).

(255, 255, 255)
track_color tuple

Màu sắc của các bản nhạc. Nếu không, các màu khác nhau sẽ được sử dụng cho các màu khác nhau Bài hát. Mặc định là Không có.

None
region_thickness int

Độ dày của các đường vùng đếm. Mặc định là 5.

5
fontsize float

Cỡ chữ cho chú thích văn bản. Mặc định là 0,7.

0.7
Mã nguồn trong ultralytics/solutions/queue_management.py
def __init__(
    self,
    classes_names,
    reg_pts=None,
    line_thickness=2,
    track_thickness=2,
    view_img=False,
    region_color=(255, 0, 255),
    view_queue_counts=True,
    draw_tracks=False,
    count_txt_color=(255, 255, 255),
    track_color=None,
    region_thickness=5,
    fontsize=0.7,
):
    """
    Initializes the QueueManager with specified parameters for tracking and counting objects.

    Args:
        classes_names (dict): A dictionary mapping class IDs to class names.
        reg_pts (list of tuples, optional): Points defining the counting region polygon. Defaults to a predefined
            rectangle.
        line_thickness (int, optional): Thickness of the annotation lines. Defaults to 2.
        track_thickness (int, optional): Thickness of the track lines. Defaults to 2.
        view_img (bool, optional): Whether to display the image frames. Defaults to False.
        region_color (tuple, optional): Color of the counting region lines (BGR). Defaults to (255, 0, 255).
        view_queue_counts (bool, optional): Whether to display the queue counts. Defaults to True.
        draw_tracks (bool, optional): Whether to draw tracks of the objects. Defaults to False.
        count_txt_color (tuple, optional): Color of the count text (BGR). Defaults to (255, 255, 255).
        track_color (tuple, optional): Color of the tracks. If None, different colors will be used for different
            tracks. Defaults to None.
        region_thickness (int, optional): Thickness of the counting region lines. Defaults to 5.
        fontsize (float, optional): Font size for the text annotations. Defaults to 0.7.
    """

    # Mouse events state
    self.is_drawing = False
    self.selected_point = None

    # Region & Line Information
    self.reg_pts = reg_pts if reg_pts is not None else [(20, 60), (20, 680), (1120, 680), (1120, 60)]
    self.counting_region = (
        Polygon(self.reg_pts) if len(self.reg_pts) >= 3 else Polygon([(20, 60), (20, 680), (1120, 680), (1120, 60)])
    )
    self.region_color = region_color
    self.region_thickness = region_thickness

    # Image and annotation Information
    self.im0 = None
    self.tf = line_thickness
    self.view_img = view_img
    self.view_queue_counts = view_queue_counts
    self.fontsize = fontsize

    self.names = classes_names  # Class names
    self.annotator = None  # Annotator
    self.window_name = "Ultralytics YOLOv8 Queue Manager"

    # Object counting Information
    self.counts = 0
    self.count_txt_color = count_txt_color

    # Tracks info
    self.track_history = defaultdict(list)
    self.track_thickness = track_thickness
    self.draw_tracks = draw_tracks
    self.track_color = track_color

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

display_frames()

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

Mã nguồn trong ultralytics/solutions/queue_management.py
def display_frames(self):
    """Displays the current frame with annotations."""
    if self.env_check:
        self.annotator.draw_region(reg_pts=self.reg_pts, thickness=self.region_thickness, color=self.region_color)
        cv2.namedWindow(self.window_name)
        cv2.imshow(self.window_name, self.im0)
        # Close window on 'q' key press
        if cv2.waitKey(1) & 0xFF == ord("q"):
            return

extract_and_process_tracks(tracks)

Trích xuất và xử lý các bản nhạc để quản lý hàng đợi trong luồng video.

Mã nguồn trong ultralytics/solutions/queue_management.py
def extract_and_process_tracks(self, tracks):
    """Extracts and processes tracks for queue management in a video stream."""

    # Initialize annotator and draw the queue region
    self.annotator = Annotator(self.im0, self.tf, self.names)

    if tracks[0].boxes.id is not None:
        boxes = tracks[0].boxes.xyxy.cpu()
        clss = tracks[0].boxes.cls.cpu().tolist()
        track_ids = tracks[0].boxes.id.int().cpu().tolist()

        # Extract tracks
        for box, track_id, cls in zip(boxes, track_ids, clss):
            # Draw bounding box
            self.annotator.box_label(box, label=f"{self.names[cls]}#{track_id}", color=colors(int(track_id), True))

            # Update track history
            track_line = self.track_history[track_id]
            track_line.append((float((box[0] + box[2]) / 2), float((box[1] + box[3]) / 2)))
            if len(track_line) > 30:
                track_line.pop(0)

            # Draw track trails if enabled
            if self.draw_tracks:
                self.annotator.draw_centroid_and_tracks(
                    track_line,
                    color=self.track_color or colors(int(track_id), True),
                    track_thickness=self.track_thickness,
                )

            prev_position = self.track_history[track_id][-2] if len(self.track_history[track_id]) > 1 else None

            # Check if the object is inside the counting region
            if len(self.reg_pts) >= 3:
                is_inside = self.counting_region.contains(Point(track_line[-1]))
                if prev_position is not None and is_inside:
                    self.counts += 1

    # Display queue counts
    label = f"Queue Counts : {str(self.counts)}"
    if label is not None:
        self.annotator.queue_counts_display(
            label,
            points=self.reg_pts,
            region_color=self.region_color,
            txt_color=self.count_txt_color,
        )

    self.counts = 0  # Reset counts after displaying
    self.display_frames()

process_queue(im0, tracks)

Chức năng chính để bắt đầu quá trình quản lý hàng đợi.

Thông số:

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

Khung hình hiện tại từ luồng video.

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
Mã nguồn trong ultralytics/solutions/queue_management.py
def process_queue(self, im0, tracks):
    """
    Main function to start the queue management process.

    Args:
        im0 (ndarray): Current frame from the video stream.
        tracks (list): List of tracks obtained from the object tracking process.
    """
    self.im0 = im0  # Store the current frame
    self.extract_and_process_tracks(tracks)  # Extract and process tracks

    if self.view_img:
        self.display_frames()  # Display the frame if enabled
    return self.im0





Created 2024-04-02, Updated 2024-06-02
Authors: glenn-jocher (2), Burhan-Q (2)