انتقل إلى المحتوى

مرجع ل ultralytics/solutions/object_counter.py

ملاحظه

هذا الملف متاح في https://github.com/ultralytics/ultralytics/ نقطة / الرئيسية /ultralytics/الحلول/object_counter.py. إذا اكتشفت مشكلة ، فيرجى المساعدة في إصلاحها من خلال المساهمة في طلب 🛠️ سحب. شكرا لك 🙏!



ultralytics.solutions.object_counter.ObjectCounter

فئة لإدارة عد الكائنات في دفق فيديو في الوقت الفعلي بناء على مساراتها.

شفرة المصدر في 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):
        """Initializes the Counter with default values for various tracking and counting parameters."""

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

        # Region & Line Information
        self.reg_pts = [(20, 400), (1260, 400)]
        self.line_dist_thresh = 15
        self.counting_region = None
        self.region_color = (255, 0, 255)
        self.region_thickness = 5

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

        self.names = None  # Classes names
        self.annotator = None  # Annotator

        # Object counting Information
        self.in_counts = 0
        self.out_counts = 0
        self.counting_list = []
        self.count_txt_thickness = 0
        self.count_txt_color = (0, 0, 0)
        self.count_color = (255, 255, 255)

        # Tracks info
        self.track_history = defaultdict(list)
        self.track_thickness = 2
        self.draw_tracks = False
        self.track_color = (0, 255, 0)

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

    def set_args(
        self,
        classes_names,
        reg_pts,
        count_reg_color=(255, 0, 255),
        line_thickness=2,
        track_thickness=2,
        view_img=False,
        view_in_counts=True,
        view_out_counts=True,
        draw_tracks=False,
        count_txt_thickness=2,
        count_txt_color=(0, 0, 0),
        count_color=(255, 255, 255),
        track_color=(0, 255, 0),
        region_thickness=5,
        line_dist_thresh=15,
    ):
        """
        Configures the Counter's image, bounding box line thickness, and counting region points.

        Args:
            line_thickness (int): Line thickness for bounding boxes.
            view_img (bool): Flag to control whether to display the video stream.
            view_in_counts (bool): Flag to control whether to display the incounts on video stream.
            view_out_counts (bool): Flag to control whether to display the outcounts on video stream.
            reg_pts (list): Initial list of points defining the counting region.
            classes_names (dict): Classes names
            track_thickness (int): Track thickness
            draw_tracks (Bool): draw tracks
            count_txt_thickness (int): Text thickness for object counting display
            count_txt_color (RGB color): count text color value
            count_color (RGB color): count text background color value
            count_reg_color (RGB color): Color of object counting region
            track_color (RGB color): color for tracks
            region_thickness (int): Object counting Region thickness
            line_dist_thresh (int): Euclidean Distance threshold for line counter
        """
        self.tf = line_thickness
        self.view_img = view_img
        self.view_in_counts = view_in_counts
        self.view_out_counts = view_out_counts
        self.track_thickness = track_thickness
        self.draw_tracks = draw_tracks

        # Region and line selection
        if len(reg_pts) == 2:
            print("Line Counter Initiated.")
            self.reg_pts = reg_pts
            self.counting_region = LineString(self.reg_pts)
        elif len(reg_pts) == 4:
            print("Region Counter Initiated.")
            self.reg_pts = reg_pts
            self.counting_region = Polygon(self.reg_pts)
        else:
            print("Invalid Region points provided, region_points can be 2 or 4")
            print("Using Line Counter Now")
            self.counting_region = LineString(self.reg_pts)

        self.names = classes_names
        self.track_color = track_color
        self.count_txt_thickness = count_txt_thickness
        self.count_txt_color = count_txt_color
        self.count_color = count_color
        self.region_color = count_reg_color
        self.region_thickness = region_thickness
        self.line_dist_thresh = line_dist_thresh

    def mouse_event_for_region(self, event, x, y, flags, params):
        """
        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.).
            params (dict): Additional parameters you may want to pass to 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."""
        boxes = tracks[0].boxes.xyxy.cpu()
        clss = tracks[0].boxes.cls.cpu().tolist()
        track_ids = tracks[0].boxes.id.int().cpu().tolist()

        # Annotator Init and region drawing
        self.annotator = Annotator(self.im0, self.tf, self.names)
        self.annotator.draw_region(reg_pts=self.reg_pts, color=self.region_color, thickness=self.region_thickness)

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

            # 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, 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
            if len(self.reg_pts) == 4:
                if (
                    prev_position is not None
                    and self.counting_region.contains(Point(track_line[-1]))
                    and track_id not in self.counting_list
                ):
                    self.counting_list.append(track_id)
                    if (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0:
                        self.in_counts += 1
                    else:
                        self.out_counts += 1

            elif len(self.reg_pts) == 2:
                if prev_position is not None:
                    distance = Point(track_line[-1]).distance(self.counting_region)
                    if distance < self.line_dist_thresh and track_id not in self.counting_list:
                        self.counting_list.append(track_id)
                        if (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0:
                            self.in_counts += 1
                        else:
                            self.out_counts += 1

        incount_label = f"In Count : {self.in_counts}"
        outcount_label = f"OutCount : {self.out_counts}"

        # Display counts based on user choice
        counts_label = None
        if not self.view_in_counts and not self.view_out_counts:
            counts_label = None
        elif not self.view_in_counts:
            counts_label = outcount_label
        elif not self.view_out_counts:
            counts_label = incount_label
        else:
            counts_label = f"{incount_label} {outcount_label}"

        if counts_label is not None:
            self.annotator.count_labels(
                counts=counts_label,
                count_txt_size=self.count_txt_thickness,
                txt_color=self.count_txt_color,
                color=self.count_color,
            )

    def display_frames(self):
        """Display frame."""
        if self.env_check:
            cv2.namedWindow("Ultralytics YOLOv8 Object Counter")
            if len(self.reg_pts) == 4:  # only add mouse event If user drawn region
                cv2.setMouseCallback(
                    "Ultralytics YOLOv8 Object Counter", self.mouse_event_for_region, {"region_points": self.reg_pts}
                )
            cv2.imshow("Ultralytics YOLOv8 Object Counter", 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

        if tracks[0].boxes.id is None:
            if self.view_img:
                self.display_frames()
            return im0
        self.extract_and_process_tracks(tracks)

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

__init__()

تهيئة العداد بقيم افتراضية لمعلمات التتبع والعد المختلفة.

شفرة المصدر في ultralytics/solutions/object_counter.py
18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 3536 37 38 3940 41 42434445 4647 484950515253545556 57
def __init__(self):
    """Initializes the Counter with default values for various tracking and counting parameters."""

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

    # Region & Line Information
    self.reg_pts = [(20, 400), (1260, 400)]
    self.line_dist_thresh = 15
    self.counting_region = None
    self.region_color = (255, 0, 255)
    self.region_thickness = 5

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

    self.names = None  # Classes names
    self.annotator = None  # Annotator

    # Object counting Information
    self.in_counts = 0
    self.out_counts = 0
    self.counting_list = []
    self.count_txt_thickness = 0
    self.count_txt_color = (0, 0, 0)
    self.count_color = (255, 255, 255)

    # Tracks info
    self.track_history = defaultdict(list)
    self.track_thickness = 2
    self.draw_tracks = False
    self.track_color = (0, 255, 0)

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

display_frames()

إطار العرض.

شفرة المصدر في ultralytics/solutions/object_counter.py
233 234 235 236 237 238239 240 241 242 243 244
def display_frames(self):
    """Display frame."""
    if self.env_check:
        cv2.namedWindow("Ultralytics YOLOv8 Object Counter")
        if len(self.reg_pts) == 4:  # only add mouse event If user drawn region
            cv2.setMouseCallback(
                "Ultralytics YOLOv8 Object Counter", self.mouse_event_for_region, {"region_points": self.reg_pts}
            )
        cv2.imshow("Ultralytics YOLOv8 Object Counter", self.im0)
        # Break Window
        if cv2.waitKey(1) & 0xFF == ord("q"):
            return

extract_and_process_tracks(tracks)

يستخرج ويعالج المسارات لحساب الكائنات في دفق الفيديو.

شفرة المصدر في ultralytics/solutions/object_counter.py
159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177178179 180 181 182 183 184 185 186 187 188189190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208209 210 211 212213214215 216 217 218219220221 222 223 224 225 226 227228229 230 231
def extract_and_process_tracks(self, tracks):
    """Extracts and processes tracks for object counting in a video stream."""
    boxes = tracks[0].boxes.xyxy.cpu()
    clss = tracks[0].boxes.cls.cpu().tolist()
    track_ids = tracks[0].boxes.id.int().cpu().tolist()

    # Annotator Init and region drawing
    self.annotator = Annotator(self.im0, self.tf, self.names)
    self.annotator.draw_region(reg_pts=self.reg_pts, color=self.region_color, thickness=self.region_thickness)

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

        # 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, 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
        if len(self.reg_pts) == 4:
            if (
                prev_position is not None
                and self.counting_region.contains(Point(track_line[-1]))
                and track_id not in self.counting_list
            ):
                self.counting_list.append(track_id)
                if (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0:
                    self.in_counts += 1
                else:
                    self.out_counts += 1

        elif len(self.reg_pts) == 2:
            if prev_position is not None:
                distance = Point(track_line[-1]).distance(self.counting_region)
                if distance < self.line_dist_thresh and track_id not in self.counting_list:
                    self.counting_list.append(track_id)
                    if (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0:
                        self.in_counts += 1
                    else:
                        self.out_counts += 1

    incount_label = f"In Count : {self.in_counts}"
    outcount_label = f"OutCount : {self.out_counts}"

    # Display counts based on user choice
    counts_label = None
    if not self.view_in_counts and not self.view_out_counts:
        counts_label = None
    elif not self.view_in_counts:
        counts_label = outcount_label
    elif not self.view_out_counts:
        counts_label = incount_label
    else:
        counts_label = f"{incount_label} {outcount_label}"

    if counts_label is not None:
        self.annotator.count_labels(
            counts=counts_label,
            count_txt_size=self.count_txt_thickness,
            txt_color=self.count_txt_color,
            color=self.count_color,
        )

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

تم تصميم هذه الوظيفة لنقل المنطقة مع أحداث الماوس في دفق فيديو في الوقت الفعلي.

البارامترات:

اسم نوع وصف افتراضي
event int

نوع حدث الماوس (على سبيل المثال، cv2. EVENT_MOUSEMOVE ، السيرة الذاتية 2. EVENT_LBUTTONDOWN ، إلخ).

مطلوب
x int

إحداثيات x لمؤشر الماوس.

مطلوب
y int

إحداثيات y لمؤشر الماوس.

مطلوب
flags int

أي علامات مرتبطة بالحدث (على سبيل المثال، cv2. EVENT_FLAG_CTRLKEY، السيرة الذاتية 2. EVENT_FLAG_SHIFTKEY ، وما إلى ذلك).

مطلوب
params dict

معلمات إضافية قد ترغب في تمريرها إلى الوظيفة.

مطلوب
شفرة المصدر في ultralytics/solutions/object_counter.py
127 128 129 130 131 132 133 134 135 136 137 138139 140 141 142 143 144 145 146147 148149150 151 152 153 154 155 156 157
def mouse_event_for_region(self, event, x, y, flags, params):
    """
    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.).
        params (dict): Additional parameters you may want to pass to 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

set_args(classes_names, reg_pts, count_reg_color=(255, 0, 255), line_thickness=2, track_thickness=2, view_img=False, view_in_counts=True, view_out_counts=True, draw_tracks=False, count_txt_thickness=2, count_txt_color=(0, 0, 0), count_color=(255, 255, 255), track_color=(0, 255, 0), region_thickness=5, line_dist_thresh=15)

تكوين صورة العداد، سمك خط المربع المحيط، وعد نقاط المنطقة.

البارامترات:

اسم نوع وصف افتراضي
line_thickness int

سمك الخط للمربعات المحيطة.

2
view_img bool

ضع علامة للتحكم في ما إذا كنت تريد عرض دفق الفيديو أم لا.

False
view_in_counts bool

ضع علامة للتحكم في ما إذا كنت تريد عرض الأعداد في دفق الفيديو.

True
view_out_counts bool

ضع علامة للتحكم في ما إذا كنت تريد عرض الأعداد الخارجة في دفق الفيديو.

True
reg_pts list

قائمة أولية بالنقاط التي تحدد منطقة العد.

مطلوب
classes_names dict

أسماء الفصول

مطلوب
track_thickness int

سمك الجنزير

2
draw_tracks Bool

رسم المسارات

False
count_txt_thickness int

سمك النص لعرض عد الكائنات

2
count_txt_color RGB color

عد قيمة لون النص

(0, 0, 0)
count_color RGB color

عد قيمة لون خلفية النص

(255, 255, 255)
count_reg_color RGB color

لون منطقة عد الكائنات

(255, 0, 255)
track_color RGB color

لون للمسارات

(0, 255, 0)
region_thickness int

عد الكائنات سمك المنطقة

5
line_dist_thresh int

عتبة المسافة الإقليدية لعداد الخط

15
شفرة المصدر في ultralytics/solutions/object_counter.py
 59 60 61 62 63 64 65 66 67 68  69 70 71 72 73 74 75 76  77 78  79 80 81 82 83     84 85 86   87 88        89   90  91   92  93    94 95 96 97 98  99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115  116 117  118 119 120 121  122123124125
def set_args(
    self,
    classes_names,
    reg_pts,
    count_reg_color=(255, 0, 255),
    line_thickness=2,
    track_thickness=2,
    view_img=False,
    view_in_counts=True,
    view_out_counts=True,
    draw_tracks=False,
    count_txt_thickness=2,
    count_txt_color=(0, 0, 0),
    count_color=(255, 255, 255),
    track_color=(0, 255, 0),
    region_thickness=5,
    line_dist_thresh=15,
):
    """
    Configures the Counter's image, bounding box line thickness, and counting region points.

    Args:
        line_thickness (int): Line thickness for bounding boxes.
        view_img (bool): Flag to control whether to display the video stream.
        view_in_counts (bool): Flag to control whether to display the incounts on video stream.
        view_out_counts (bool): Flag to control whether to display the outcounts on video stream.
        reg_pts (list): Initial list of points defining the counting region.
        classes_names (dict): Classes names
        track_thickness (int): Track thickness
        draw_tracks (Bool): draw tracks
        count_txt_thickness (int): Text thickness for object counting display
        count_txt_color (RGB color): count text color value
        count_color (RGB color): count text background color value
        count_reg_color (RGB color): Color of object counting region
        track_color (RGB color): color for tracks
        region_thickness (int): Object counting Region thickness
        line_dist_thresh (int): Euclidean Distance threshold for line counter
    """
    self.tf = line_thickness
    self.view_img = view_img
    self.view_in_counts = view_in_counts
    self.view_out_counts = view_out_counts
    self.track_thickness = track_thickness
    self.draw_tracks = draw_tracks

    # Region and line selection
    if len(reg_pts) == 2:
        print("Line Counter Initiated.")
        self.reg_pts = reg_pts
        self.counting_region = LineString(self.reg_pts)
    elif len(reg_pts) == 4:
        print("Region Counter Initiated.")
        self.reg_pts = reg_pts
        self.counting_region = Polygon(self.reg_pts)
    else:
        print("Invalid Region points provided, region_points can be 2 or 4")
        print("Using Line Counter Now")
        self.counting_region = LineString(self.reg_pts)

    self.names = classes_names
    self.track_color = track_color
    self.count_txt_thickness = count_txt_thickness
    self.count_txt_color = count_txt_color
    self.count_color = count_color
    self.region_color = count_reg_color
    self.region_thickness = region_thickness
    self.line_dist_thresh = line_dist_thresh

start_counting(im0, tracks)

الوظيفة الرئيسية لبدء عملية عد الكائنات.

البارامترات:

اسم نوع وصف افتراضي
im0 ndarray

الإطار الحالي من دفق الفيديو.

مطلوب
tracks list

قائمة المسارات التي تم الحصول عليها من عملية تتبع الكائن.

مطلوب
شفرة المصدر في ultralytics/solutions/object_counter.py
 246 247 248 249 250 251 252 253 254 255 256 257258259 260261 262 263264
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

    if tracks[0].boxes.id is None:
        if self.view_img:
            self.display_frames()
        return im0
    self.extract_and_process_tracks(tracks)

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





تم إنشاء 2023-12-02, اخر تحديث 2023-12-02
المؤلفون: رضوان منور (1)