Spaces:
Runtime error
Runtime error
# Ultralytics YOLO π, AGPL-3.0 license | |
from collections import defaultdict | |
import logging | |
import cv2 | |
import numpy as np | |
from ultralytics.utils.checks import check_imshow, check_requirements | |
from ultralytics.utils.plotting import Annotator, colors | |
from shapely.geometry import LineString, Point, Polygon | |
# create logger | |
logger = logging.getLogger(__name__).addHandler(logging.NullHandler()) | |
# need shapely>=2.0.0 | |
check_requirements("shapely>=2.0.0") | |
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.out_counts_prev = self.out_counts | |
self.in_counts_prev = self.in_counts | |
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.draw_boxes = False # added by steve.yin @ 3/1/2024 | |
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, | |
draw_boxes=False, # added by steve.yin @ 3/1/2024 | |
draw_reg_pts=True, # added by steve.yin @ 3/1/2024 | |
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 display the video stream. | |
view_in_counts (bool): Flag to control display the incounts. | |
view_out_counts (bool): Flag to control display the outcounts. | |
reg_pts (list): Initial list of points for the counting region. | |
classes_names (dict): Classes names | |
track_thickness (int): Track thickness | |
draw_tracks (Bool): draw tracks | |
draw_boxes (Bool): draw boxes | |
draw_reg_pts (Bool): draw reg_pts | |
count_txt_thickness (int): Text thickness 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 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 | |
self.draw_boxes = draw_boxes # added by steve.yin @ 3/1/2024 | |
self.draw_reg_pts = draw_reg_pts # added by steve.yin @ 3/1/2024 | |
# Region and line selection | |
if len(reg_pts) == 2: | |
logging.info("Line Counter Initiated.") | |
self.reg_pts = reg_pts | |
self.counting_region = LineString(self.reg_pts) | |
u = np.array([self.reg_pts[0][0], self.reg_pts[0][1]]) | |
v = np.array([self.reg_pts[1][0], self.reg_pts[1][1]]) | |
elif len(reg_pts) == 4: | |
logging.info("Region Counter Initiated.") | |
self.reg_pts = reg_pts | |
self.counting_region = Polygon(self.reg_pts) | |
u = np.array([ | |
(self.reg_pts[0][0] + self.reg_pts[1][0]) / 2, | |
(self.reg_pts[0][1] + self.reg_pts[1][1]) / 2, | |
]) | |
v = np.array([ | |
(self.reg_pts[2][0] + self.reg_pts[3][0]) / 2, | |
(self.reg_pts[2][1] + self.reg_pts[3][1]) / 2, | |
]) | |
else: | |
logging.warning( | |
"Invalid Region points, which can only be 2 or 4. " + | |
"Using Line Counter Instead!" | |
) | |
self.counting_region = LineString(self.reg_pts) | |
u = np.array(self.counting_region.coords[0]) | |
v = np.array( | |
self.counting_region.coords[len(self.counting_region.coords)-1] | |
) | |
# get line orientation, rotate ccw 90degrees, get line normal vector | |
n = v - u | |
nvec = np.array([-n[1], n[0]]) | |
# print(f"v: {v}, u: {u}, n: {n}, nvec0: {nvec}") | |
self.counting_region_nvec = nvec / (np.linalg.norm(nvec) + 1e-6) | |
# print(f"nvec: {self.counting_region_nvec}") | |
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 passing 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 [modified by steve.yin @ 3/1/2024] | |
if self.draw_reg_pts: | |
self.annotator.draw_region( | |
reg_pts=self.reg_pts, | |
color=self.region_color, | |
thickness=self.region_thickness | |
) | |
if self.draw_boxes: | |
self.annotator.box_label( | |
box=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=track_line, | |
color=self.track_color, | |
track_thickness=self.track_thickness | |
) | |
prev_position = self.track_history[track_id][0] \ | |
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) | |
obj_track_vec = np.array([ | |
track_line[-1][0] - prev_position[0], | |
track_line[-1][1] - prev_position[1] | |
]) | |
if np.sign( | |
np.dot(obj_track_vec, self.counting_region_nvec) | |
) < 0: | |
self.out_counts += 1 | |
else: | |
self.in_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) | |
obj_track_vec = np.array([ | |
track_line[-1][0] - prev_position[0], | |
track_line[-1][1] - prev_position[1] | |
]) | |
logging.info(f"obj_track_vec: {obj_track_vec}") | |
if np.sign( | |
np.dot(obj_track_vec, self.counting_region_nvec) | |
) < 0: | |
self.out_counts += 1 | |
else: | |
self.in_counts += 1 | |
self.outcounts_updated() | |
self.incounts_updated() | |
self.out_counts_prev = self.out_counts | |
self.in_counts_prev = self.in_counts | |
incount_label = f"In: {self.in_counts}" | |
outcount_label = f"Out: {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") | |
# only add mouse event If user drawn region | |
if len(self.reg_pts) == 4: | |
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 | |
def incounts_updated(self): | |
if self.in_counts_prev < self.in_counts: | |
yield f"{self.in_counts}" | |
def outcounts_updated(self): | |
if self.out_counts_prev < self.out_counts: | |
yield f"{self.out_counts}" | |
if __name__ == "__main__": | |
ObjectCounter() | |