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

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

Ghi

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

Một lớp để quản lý việc đếm các đố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/object_counter.py
class ObjectCounter:
    """A class to manage the counting of objects in a real-time video stream based on their tracks."""

    def __init__(
        self,
        classes_names,
        reg_pts=None,
        count_reg_color=(255, 0, 255),
        count_txt_color=(0, 0, 0),
        count_bg_color=(255, 255, 255),
        line_thickness=2,
        track_thickness=2,
        view_img=False,
        view_in_counts=True,
        view_out_counts=True,
        draw_tracks=False,
        track_color=None,
        region_thickness=5,
        line_dist_thresh=15,
        cls_txtdisplay_gap=50,
    ):
        """
        Initializes the ObjectCounter with various tracking and counting parameters.

        Args:
            classes_names (dict): Dictionary of class names.
            reg_pts (list): List of points defining the counting region.
            count_reg_color (tuple): RGB color of the counting region.
            count_txt_color (tuple): RGB color of the count text.
            count_bg_color (tuple): RGB color of the count text background.
            line_thickness (int): Line thickness for bounding boxes.
            track_thickness (int): Thickness of the track lines.
            view_img (bool): Flag to control whether to display the video stream.
            view_in_counts (bool): Flag to control whether to display the in counts on the video stream.
            view_out_counts (bool): Flag to control whether to display the out counts on the video stream.
            draw_tracks (bool): Flag to control whether to draw the object tracks.
            track_color (tuple): RGB color of the tracks.
            region_thickness (int): Thickness of the object counting region.
            line_dist_thresh (int): Euclidean distance threshold for line counter.
            cls_txtdisplay_gap (int): Display gap between each class count.
        """

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

        # Region & Line Information
        self.reg_pts = [(20, 400), (1260, 400)] if reg_pts is None else reg_pts
        self.line_dist_thresh = line_dist_thresh
        self.counting_region = None
        self.region_color = count_reg_color
        self.region_thickness = region_thickness

        # Image and annotation Information
        self.im0 = None
        self.tf = line_thickness
        self.view_img = view_img
        self.view_in_counts = view_in_counts
        self.view_out_counts = view_out_counts

        self.names = classes_names  # Classes names
        self.annotator = None  # Annotator
        self.window_name = "Ultralytics YOLOv8 Object Counter"

        # Object counting Information
        self.in_counts = 0
        self.out_counts = 0
        self.count_ids = []
        self.class_wise_count = {}
        self.count_txt_thickness = 0
        self.count_txt_color = count_txt_color
        self.count_bg_color = count_bg_color
        self.cls_txtdisplay_gap = cls_txtdisplay_gap
        self.fontsize = 0.6

        # 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)

        # Initialize counting region
        if len(self.reg_pts) == 2:
            print("Line Counter Initiated.")
            self.counting_region = LineString(self.reg_pts)
        elif len(self.reg_pts) >= 3:
            print("Polygon Counter Initiated.")
            self.counting_region = Polygon(self.reg_pts)
        else:
            print("Invalid Region points provided, region_points must be 2 for lines or >= 3 for polygons.")
            print("Using Line Counter Now")
            self.counting_region = LineString(self.reg_pts)

    def mouse_event_for_region(self, event, x, y, flags, params):
        """
        Handles mouse events for defining and moving the counting region 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 associated event flags (e.g., cv2.EVENT_FLAG_CTRLKEY,  cv2.EVENT_FLAG_SHIFTKEY, etc.).
            params (dict): Additional parameters for the function.
        """
        if event == cv2.EVENT_LBUTTONDOWN:
            for i, point in enumerate(self.reg_pts):
                if (
                    isinstance(point, (tuple, list))
                    and len(point) >= 2
                    and (abs(x - point[0]) < 10 and abs(y - point[1]) < 10)
                ):
                    self.selected_point = i
                    self.is_drawing = True
                    break

        elif event == cv2.EVENT_MOUSEMOVE:
            if self.is_drawing and self.selected_point is not None:
                self.reg_pts[self.selected_point] = (x, y)
                self.counting_region = Polygon(self.reg_pts)

        elif event == cv2.EVENT_LBUTTONUP:
            self.is_drawing = False
            self.selected_point = None

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

        # Annotator Init and region drawing
        self.annotator = Annotator(self.im0, self.tf, self.names)

        # Draw region or line
        self.annotator.draw_region(reg_pts=self.reg_pts, color=self.region_color, thickness=self.region_thickness)

        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))

                # Store class info
                if self.names[cls] not in self.class_wise_count:
                    self.class_wise_count[self.names[cls]] = {"IN": 0, "OUT": 0}

                # Draw Tracks
                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 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

                # Count objects in any polygon
                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 and track_id not in self.count_ids:
                        self.count_ids.append(track_id)

                        if (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0:
                            self.in_counts += 1
                            self.class_wise_count[self.names[cls]]["IN"] += 1
                        else:
                            self.out_counts += 1
                            self.class_wise_count[self.names[cls]]["OUT"] += 1

                # Count objects using line
                elif len(self.reg_pts) == 2:
                    if prev_position is not None and track_id not in self.count_ids:
                        distance = Point(track_line[-1]).distance(self.counting_region)
                        if distance < self.line_dist_thresh and track_id not in self.count_ids:
                            self.count_ids.append(track_id)

                            if (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0:
                                self.in_counts += 1
                                self.class_wise_count[self.names[cls]]["IN"] += 1
                            else:
                                self.out_counts += 1
                                self.class_wise_count[self.names[cls]]["OUT"] += 1

        labels_dict = {}

        for key, value in self.class_wise_count.items():
            if value["IN"] != 0 or value["OUT"] != 0:
                if not self.view_in_counts and not self.view_out_counts:
                    continue
                elif not self.view_in_counts:
                    labels_dict[str.capitalize(key)] = f"OUT {value['OUT']}"
                elif not self.view_out_counts:
                    labels_dict[str.capitalize(key)] = f"IN {value['IN']}"
                else:
                    labels_dict[str.capitalize(key)] = f"IN {value['IN']} OUT {value['OUT']}"

        if labels_dict:
            self.annotator.display_analytics(self.im0, labels_dict, self.count_txt_color, self.count_bg_color, 10)

    def display_frames(self):
        """Displays the current frame with annotations and regions in a window."""
        if self.env_check:
            cv2.namedWindow(self.window_name)
            if len(self.reg_pts) == 4:  # only add mouse event If user drawn region
                cv2.setMouseCallback(self.window_name, self.mouse_event_for_region, {"region_points": self.reg_pts})
            cv2.imshow(self.window_name, self.im0)
            # Break Window
            if cv2.waitKey(1) & 0xFF == ord("q"):
                return

    def start_counting(self, im0, tracks):
        """
        Main function to start the object counting 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 image
        self.extract_and_process_tracks(tracks)  # draw region even if no objects

        if self.view_img:
            self.display_frames()
        return self.im0

__init__(classes_names, reg_pts=None, count_reg_color=(255, 0, 255), count_txt_color=(0, 0, 0), count_bg_color=(255, 255, 255), line_thickness=2, track_thickness=2, view_img=False, view_in_counts=True, view_out_counts=True, draw_tracks=False, track_color=None, region_thickness=5, line_dist_thresh=15, cls_txtdisplay_gap=50)

Khởi tạo ObjectCounter với các tham số theo dõi và đếm khác nhau.

Thông số:

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

Từ điển tên lớp.

bắt buộc
reg_pts list

Danh sách các điểm xác định khu vực đếm.

None
count_reg_color tuple

Màu RGB của vùng đếm.

(255, 0, 255)
count_txt_color tuple

Màu RGB của văn bản đếm.

(0, 0, 0)
count_bg_color tuple

Màu RGB của nền văn bản đếm.

(255, 255, 255)
line_thickness int

Độ dày dòng cho các hộp giới hạn.

2
track_thickness int

Độ dày của các đường ray.

2
view_img bool

Gắn cờ để kiểm soát việc có hiển thị luồng video hay không.

False
view_in_counts bool

Gắn cờ để kiểm soát việc có hiển thị số vào trên luồng video hay không.

True
view_out_counts bool

Gắn cờ để kiểm soát việc có hiển thị số lượng ra trên luồng video hay không.

True
draw_tracks bool

Gắn cờ để kiểm soát xem có vẽ các rãnh đối tượng hay không.

False
track_color tuple

Màu RGB của các bản nhạc.

None
region_thickness int

Độ dày của vùng đếm đối tượng.

5
line_dist_thresh int

Ngưỡng khoảng cách Euclide cho bộ đếm dòng.

15
cls_txtdisplay_gap int

Hiển thị khoảng cách giữa mỗi lớp học.

50
Mã nguồn trong ultralytics/solutions/object_counter.py
def __init__(
    self,
    classes_names,
    reg_pts=None,
    count_reg_color=(255, 0, 255),
    count_txt_color=(0, 0, 0),
    count_bg_color=(255, 255, 255),
    line_thickness=2,
    track_thickness=2,
    view_img=False,
    view_in_counts=True,
    view_out_counts=True,
    draw_tracks=False,
    track_color=None,
    region_thickness=5,
    line_dist_thresh=15,
    cls_txtdisplay_gap=50,
):
    """
    Initializes the ObjectCounter with various tracking and counting parameters.

    Args:
        classes_names (dict): Dictionary of class names.
        reg_pts (list): List of points defining the counting region.
        count_reg_color (tuple): RGB color of the counting region.
        count_txt_color (tuple): RGB color of the count text.
        count_bg_color (tuple): RGB color of the count text background.
        line_thickness (int): Line thickness for bounding boxes.
        track_thickness (int): Thickness of the track lines.
        view_img (bool): Flag to control whether to display the video stream.
        view_in_counts (bool): Flag to control whether to display the in counts on the video stream.
        view_out_counts (bool): Flag to control whether to display the out counts on the video stream.
        draw_tracks (bool): Flag to control whether to draw the object tracks.
        track_color (tuple): RGB color of the tracks.
        region_thickness (int): Thickness of the object counting region.
        line_dist_thresh (int): Euclidean distance threshold for line counter.
        cls_txtdisplay_gap (int): Display gap between each class count.
    """

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

    # Region & Line Information
    self.reg_pts = [(20, 400), (1260, 400)] if reg_pts is None else reg_pts
    self.line_dist_thresh = line_dist_thresh
    self.counting_region = None
    self.region_color = count_reg_color
    self.region_thickness = region_thickness

    # Image and annotation Information
    self.im0 = None
    self.tf = line_thickness
    self.view_img = view_img
    self.view_in_counts = view_in_counts
    self.view_out_counts = view_out_counts

    self.names = classes_names  # Classes names
    self.annotator = None  # Annotator
    self.window_name = "Ultralytics YOLOv8 Object Counter"

    # Object counting Information
    self.in_counts = 0
    self.out_counts = 0
    self.count_ids = []
    self.class_wise_count = {}
    self.count_txt_thickness = 0
    self.count_txt_color = count_txt_color
    self.count_bg_color = count_bg_color
    self.cls_txtdisplay_gap = cls_txtdisplay_gap
    self.fontsize = 0.6

    # 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)

    # Initialize counting region
    if len(self.reg_pts) == 2:
        print("Line Counter Initiated.")
        self.counting_region = LineString(self.reg_pts)
    elif len(self.reg_pts) >= 3:
        print("Polygon Counter Initiated.")
        self.counting_region = Polygon(self.reg_pts)
    else:
        print("Invalid Region points provided, region_points must be 2 for lines or >= 3 for polygons.")
        print("Using Line Counter Now")
        self.counting_region = LineString(self.reg_pts)

display_frames()

Hiển thị khung hiện tại với chú thích và vùng trong cửa sổ.

Mã nguồn trong ultralytics/solutions/object_counter.py
def display_frames(self):
    """Displays the current frame with annotations and regions in a window."""
    if self.env_check:
        cv2.namedWindow(self.window_name)
        if len(self.reg_pts) == 4:  # only add mouse event If user drawn region
            cv2.setMouseCallback(self.window_name, self.mouse_event_for_region, {"region_points": self.reg_pts})
        cv2.imshow(self.window_name, self.im0)
        # Break Window
        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 để đếm đối tượng trong luồng video.

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

    # Annotator Init and region drawing
    self.annotator = Annotator(self.im0, self.tf, self.names)

    # Draw region or line
    self.annotator.draw_region(reg_pts=self.reg_pts, color=self.region_color, thickness=self.region_thickness)

    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))

            # Store class info
            if self.names[cls] not in self.class_wise_count:
                self.class_wise_count[self.names[cls]] = {"IN": 0, "OUT": 0}

            # Draw Tracks
            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 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

            # Count objects in any polygon
            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 and track_id not in self.count_ids:
                    self.count_ids.append(track_id)

                    if (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0:
                        self.in_counts += 1
                        self.class_wise_count[self.names[cls]]["IN"] += 1
                    else:
                        self.out_counts += 1
                        self.class_wise_count[self.names[cls]]["OUT"] += 1

            # Count objects using line
            elif len(self.reg_pts) == 2:
                if prev_position is not None and track_id not in self.count_ids:
                    distance = Point(track_line[-1]).distance(self.counting_region)
                    if distance < self.line_dist_thresh and track_id not in self.count_ids:
                        self.count_ids.append(track_id)

                        if (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0:
                            self.in_counts += 1
                            self.class_wise_count[self.names[cls]]["IN"] += 1
                        else:
                            self.out_counts += 1
                            self.class_wise_count[self.names[cls]]["OUT"] += 1

    labels_dict = {}

    for key, value in self.class_wise_count.items():
        if value["IN"] != 0 or value["OUT"] != 0:
            if not self.view_in_counts and not self.view_out_counts:
                continue
            elif not self.view_in_counts:
                labels_dict[str.capitalize(key)] = f"OUT {value['OUT']}"
            elif not self.view_out_counts:
                labels_dict[str.capitalize(key)] = f"IN {value['IN']}"
            else:
                labels_dict[str.capitalize(key)] = f"IN {value['IN']} OUT {value['OUT']}"

    if labels_dict:
        self.annotator.display_analytics(self.im0, labels_dict, self.count_txt_color, self.count_bg_color, 10)

mouse_event_for_region(event, x, y, flags, params)

Xử lý các sự kiện chuột để xác định và di chuyển vùng đếm 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

Mọi cờ sự kiện liên quan (ví dụ: cv2. EVENT_FLAG_CTRLKEY, CV2. EVENT_FLAG_SHIFTKEY, v.v.).

bắt buộc
params dict

Các tham số bổ sung cho hàm.

bắt buộc
Mã nguồn trong ultralytics/solutions/object_counter.py
def mouse_event_for_region(self, event, x, y, flags, params):
    """
    Handles mouse events for defining and moving the counting region 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 associated event flags (e.g., cv2.EVENT_FLAG_CTRLKEY,  cv2.EVENT_FLAG_SHIFTKEY, etc.).
        params (dict): Additional parameters for the function.
    """
    if event == cv2.EVENT_LBUTTONDOWN:
        for i, point in enumerate(self.reg_pts):
            if (
                isinstance(point, (tuple, list))
                and len(point) >= 2
                and (abs(x - point[0]) < 10 and abs(y - point[1]) < 10)
            ):
                self.selected_point = i
                self.is_drawing = True
                break

    elif event == cv2.EVENT_MOUSEMOVE:
        if self.is_drawing and self.selected_point is not None:
            self.reg_pts[self.selected_point] = (x, y)
            self.counting_region = Polygon(self.reg_pts)

    elif event == cv2.EVENT_LBUTTONUP:
        self.is_drawing = False
        self.selected_point = None

start_counting(im0, tracks)

Chức năng chính để bắt đầu quá trình đếm đối tượng.

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/object_counter.py
def start_counting(self, im0, tracks):
    """
    Main function to start the object counting 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 image
    self.extract_and_process_tracks(tracks)  # draw region even if no objects

    if self.view_img:
        self.display_frames()
    return self.im0





Created 2023-12-02, Updated 2024-06-02
Authors: glenn-jocher (2), Burhan-Q (1), RizwanMunawar (1)