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.

Source code in ultralytics/solutions/speed_estimation.py 
14
15
16
17
18
19
20
21
22
23
def __init__(self, **kwargs):
    """Initializes the SpeedEstimator with the given parameters."""
    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:

Name Type Description Default
im0 ndarray

The input image that will be used for processing

 required

Returns im0 (ndarray): The processed image for more usage

Source code in ultralytics/solutions/speed_estimation.py 
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
def estimate_speed(self, im0):
    """
    Estimates the speed of objects based on tracking data.

    Args:
        im0 (ndarray): The input image that will be used for processing
    Returns
        im0 (ndarray): The processed image for more usage
    """
    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 LineString([self.trk_pp[track_id], self.track_line[-1]]).intersects(self.l_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 9 months ago ✏️ Updated 1 month ago
jk4e glenn-jocher Burhan-Q AyushExel RizwanMunawar