Spaces:
Running
Running
File size: 4,959 Bytes
1999a98 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 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 |
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from ultralytics.solutions.solutions import BaseSolution
from ultralytics.utils.plotting import Annotator, colors
class QueueManager(BaseSolution):
"""
Manages queue counting in real-time video streams based on object tracks.
This class extends BaseSolution to provide functionality for tracking and counting objects within a specified
region in video frames.
Attributes:
counts (int): The current count of objects in the queue.
rect_color (Tuple[int, int, int]): RGB color tuple for drawing the queue region rectangle.
region_length (int): The number of points defining the queue region.
annotator (Annotator): An instance of the Annotator class for drawing on frames.
track_line (List[Tuple[int, int]]): List of track line coordinates.
track_history (Dict[int, List[Tuple[int, int]]]): Dictionary storing tracking history for each object.
Methods:
initialize_region: Initializes the queue region.
process_queue: Processes a single frame for queue management.
extract_tracks: Extracts object tracks from the current frame.
store_tracking_history: Stores the tracking history for an object.
display_output: Displays the processed output.
Examples:
>>> cap = cv2.VideoCapture("Path/to/video/file.mp4")
>>> queue_manager = QueueManager(region=[100, 100, 200, 200, 300, 300])
>>> while cap.isOpened():
>>> success, im0 = cap.read()
>>> if not success:
>>> break
>>> out = queue.process_queue(im0)
"""
def __init__(self, **kwargs):
"""Initializes the QueueManager with parameters for tracking and counting objects in a video stream."""
super().__init__(**kwargs)
self.initialize_region()
self.counts = 0 # Queue counts Information
self.rect_color = (255, 255, 255) # Rectangle color
self.region_length = len(self.region) # Store region length for further usage
def process_queue(self, im0):
"""
Processes the queue management for a single frame of video.
Args:
im0 (numpy.ndarray): Input image for processing, typically a frame from a video stream.
Returns:
(numpy.ndarray): Processed image with annotations, bounding boxes, and queue counts.
This method performs the following steps:
1. Resets the queue count for the current frame.
2. Initializes an Annotator object for drawing on the image.
3. Extracts tracks from the image.
4. Draws the counting region on the image.
5. For each detected object:
- Draws bounding boxes and labels.
- Stores tracking history.
- Draws centroids and tracks.
- Checks if the object is inside the counting region and updates the count.
6. Displays the queue count on the image.
7. Displays the processed output.
Examples:
>>> queue_manager = QueueManager()
>>> frame = cv2.imread("frame.jpg")
>>> processed_frame = queue_manager.process_queue(frame)
"""
self.counts = 0 # Reset counts every frame
self.annotator = Annotator(im0, line_width=self.line_width) # Initialize annotator
self.extract_tracks(im0) # Extract tracks
self.annotator.draw_region(
reg_pts=self.region, color=self.rect_color, thickness=self.line_width * 2
) # Draw region
for box, track_id, cls in zip(self.boxes, self.track_ids, self.clss):
# Draw bounding box and counting region
self.annotator.box_label(box, label=self.names[cls], color=colors(track_id, True))
self.store_tracking_history(track_id, box) # Store track history
# Draw tracks of objects
self.annotator.draw_centroid_and_tracks(
self.track_line, color=colors(int(track_id), True), track_thickness=self.line_width
)
# Cache frequently accessed attributes
track_history = self.track_history.get(track_id, [])
# store previous position of track and check if the object is inside the counting region
prev_position = None
if len(track_history) > 1:
prev_position = track_history[-2]
if self.region_length >= 3 and prev_position and self.r_s.contains(self.Point(self.track_line[-1])):
self.counts += 1
# Display queue counts
self.annotator.queue_counts_display(
f"Queue Counts : {str(self.counts)}",
points=self.region,
region_color=self.rect_color,
txt_color=(104, 31, 17),
)
self.display_output(im0) # display output with base class function
return im0 # return output image for more usage
|