Saltar al contenido

Referencia para ultralytics/solutions/object_counter.py

Nota

Este archivo est谩 disponible en https://github.com/ultralytics/ ultralytics/blob/main/ ultralytics/solutions/object_counter .py. Si detectas alg煤n problema, por favor, ayuda a solucionarlo contribuyendo con una Pull Request 馃洜锔. 隆Gracias 馃檹!



ultralytics.solutions.object_counter.ObjectCounter

Una clase para gestionar el recuento de objetos en un flujo de v铆deo en tiempo real bas谩ndose en sus huellas.

C贸digo fuente en 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 if self.track_color else 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)

Inicializa el ObjectCounter con varios par谩metros de seguimiento y recuento.

Par谩metros:

Nombre Tipo Descripci贸n Por defecto
classes_names dict

Diccionario de nombres de clase.

necesario
reg_pts list

Lista de puntos que definen la regi贸n de recuento.

None
count_reg_color tuple

Color RGB de la regi贸n de recuento.

(255, 0, 255)
count_txt_color tuple

Color RGB del texto de recuento.

(0, 0, 0)
count_bg_color tuple

Color RGB del fondo del texto de recuento.

(255, 255, 255)
line_thickness int

Grosor de l铆nea de los cuadros delimitadores.

2
track_thickness int

Grosor de las l铆neas de la v铆a.

2
view_img bool

Bandera para controlar si se muestra el flujo de v铆deo.

False
view_in_counts bool

Indicador para controlar si se muestran los recuentos en el flujo de v铆deo.

True
view_out_counts bool

Bandera para controlar si se muestran los recuentos de salida en el flujo de v铆deo.

True
draw_tracks bool

Indicador para controlar si se dibujan las pistas del objeto.

False
track_color tuple

Color RGB de las pistas.

None
region_thickness int

Grosor de la regi贸n de recuento de objetos.

5
line_dist_thresh int

Umbral de distancia euclidiana para el contador de l铆neas.

15
cls_txtdisplay_gap int

Muestra el espacio entre cada recuento de clases.

50
C贸digo fuente en 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()

Muestra el fotograma actual con anotaciones y regiones en una ventana.

C贸digo fuente en 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)

Extrae y procesa pistas para el recuento de objetos en un flujo de v铆deo.

C贸digo fuente en 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 if self.track_color else 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)

Maneja los eventos del rat贸n para definir y mover la regi贸n de recuento en un flujo de v铆deo en tiempo real.

Par谩metros:

Nombre Tipo Descripci贸n Por defecto
event int

El tipo de evento del rat贸n (por ejemplo, cv2.EVENT_MOUSEMOVE, cv2.EVENT_LBUTTONDOWN, etc.).

necesario
x int

La coordenada x del puntero del rat贸n.

necesario
y int

La coordenada y del puntero del rat贸n.

necesario
flags int

Cualquier indicador de evento asociado (por ejemplo, cv2.EVENT_FLAG_CTRLKEY, cv2.EVENT_FLAG_SHIFTKEY, etc.).

necesario
params dict

Par谩metros adicionales para la funci贸n.

necesario
C贸digo fuente en 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)

Funci贸n principal para iniciar el proceso de recuento de objetos.

Par谩metros:

Nombre Tipo Descripci贸n Por defecto
im0 ndarray

Fotograma actual del flujo de v铆deo.

necesario
tracks list

Lista de pistas obtenidas en el proceso de seguimiento de objetos.

necesario
C贸digo fuente en 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





Creado 2023-12-02, Actualizado 2024-05-18
Autores: glenn-jocher (1), Burhan-Q (1), RizwanMunawar (1)