Skip to content

Reference for ultralytics/solutions/speed_estimation.py

Note

This file is available at https://github.com/ultralytics/ultralytics/blob/main/ultralytics/solutions/speed_estimation.py. If you spot a problem please help fix it by contributing a Pull Request 🛠️. Thank you 🙏!


ultralytics.solutions.speed_estimation.SpeedEstimator

SpeedEstimator(**kwargs)

Bases: BaseSolution

A class to estimate the speed of objects in a real-time video stream based on their tracks.

This class extends the BaseSolution class and provides functionality for estimating object speeds using tracking data in video streams.

Attributes:

NameTypeDescription
spdDict[int, float]

Dictionary storing speed data for tracked objects.

trkd_idsList[int]

List of tracked object IDs that have already been speed-estimated.

trk_ptDict[int, float]

Dictionary storing previous timestamps for tracked objects.

trk_ppDict[int, Tuple[float, float]]

Dictionary storing previous positions for tracked objects.

annotatorAnnotator

Annotator object for drawing on images.

regionList[Tuple[int, int]]

List of points defining the speed estimation region.

track_lineList[Tuple[float, float]]

List of points representing the object's track.

r_sLineString

LineString object representing the speed estimation region.

Methods:

NameDescription
initialize_region

Initializes the speed estimation region.

estimate_speed

Estimates the speed of objects based on tracking data.

store_tracking_history

Stores the tracking history for an object.

extract_tracks

Extracts tracks from the current frame.

display_output

Displays the output with annotations.

Examples:

>>> estimator = SpeedEstimator()
>>> frame = cv2.imread("frame.jpg")
>>> processed_frame = estimator.estimate_speed(frame)
>>> cv2.imshow("Speed Estimation", processed_frame)
Source code in ultralytics/solutions/speed_estimation.py 
42
43
44
45
46
47
48
49
50
51
def __init__(self, **kwargs):
    """Initializes the SpeedEstimator object with speed estimation parameters and data structures."""
    super().__init__(**kwargs)

    self.initialize_region()  # Initialize speed region

    self.spd = {}  # set for speed data
    self.trkd_ids = []  # list for already speed_estimated and tracked ID's
    self.trk_pt = {}  # set for tracks previous time
    self.trk_pp = {}  # set for tracks previous point

estimate_speed

estimate_speed(im0)

Estimates the speed of objects based on tracking data.

Parameters:

NameTypeDescriptionDefault
im0ndarray

Input image for processing. Shape is typically (H, W, C) for RGB images.

required

Returns:

TypeDescription
ndarray

Processed image with speed estimations and annotations.

Examples:

>>> estimator = SpeedEstimator()
>>> image = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
>>> processed_image = estimator.estimate_speed(image)
Source code in ultralytics/solutions/speed_estimation.py 
 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
def estimate_speed(self, im0):
    """
    Estimates the speed of objects based on tracking data.

    Args:
        im0 (np.ndarray): Input image for processing. Shape is typically (H, W, C) for RGB images.

    Returns:
        (np.ndarray): Processed image with speed estimations and annotations.

    Examples:
        >>> estimator = SpeedEstimator()
        >>> image = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
        >>> processed_image = estimator.estimate_speed(image)
    """
    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=(104, 0, 123), thickness=self.line_width * 2
    )  # Draw region

    for box, track_id, cls in zip(self.boxes, self.track_ids, self.clss):
        self.store_tracking_history(track_id, box)  # Store track history

        # Check if track_id is already in self.trk_pp or trk_pt initialize if not
        if track_id not in self.trk_pt:
            self.trk_pt[track_id] = 0
        if track_id not in self.trk_pp:
            self.trk_pp[track_id] = self.track_line[-1]

        speed_label = f"{int(self.spd[track_id])} km/h" if track_id in self.spd else self.names[int(cls)]
        self.annotator.box_label(box, label=speed_label, color=colors(track_id, True))  # Draw bounding box

        # Draw tracks of objects
        self.annotator.draw_centroid_and_tracks(
            self.track_line, color=colors(int(track_id), True), track_thickness=self.line_width
        )

        # Calculate object speed and direction based on region intersection
        if self.LineString([self.trk_pp[track_id], self.track_line[-1]]).intersects(self.r_s):
            direction = "known"
        else:
            direction = "unknown"

        # Perform speed calculation and tracking updates if direction is valid
        if direction == "known" and track_id not in self.trkd_ids:
            self.trkd_ids.append(track_id)
            time_difference = time() - self.trk_pt[track_id]
            if time_difference > 0:
                self.spd[track_id] = np.abs(self.track_line[-1][1] - self.trk_pp[track_id][1]) / time_difference

        self.trk_pt[track_id] = time()
        self.trk_pp[track_id] = self.track_line[-1]

    self.display_output(im0)  # display output with base class function

    return im0  # return output image for more usage



📅 Created 10 months ago ✏️ Updated 2 months ago
jk4e glenn-jocher Burhan-Q AyushExel RizwanMunawar