跳至内容

参考资料 ultralytics/solutions/speed_estimation.py

备注

该文件可在https://github.com/ultralytics/ultralytics/blob/main/ultralytics/solutions/speed_estimation .py 下找到。如果您发现问题,请通过提交 Pull Request🛠️ 帮助修复。谢谢🙏!



ultralytics.solutions.speed_estimation.SpeedEstimator

根据实时视频流中物体的轨迹估算其速度的类。

源代码 ultralytics/solutions/speed_estimation.py
class SpeedEstimator:
    """A class to estimation speed of objects in real-time video stream based on their tracks."""

    def __init__(self):
        """Initializes the speed-estimator class with default values for Visual, Image, track and speed parameters."""

        # Visual & im0 information
        self.im0 = None
        self.annotator = None
        self.view_img = False

        # Region information
        self.reg_pts = [(20, 400), (1260, 400)]
        self.region_thickness = 3

        # Predict/track information
        self.clss = None
        self.names = None
        self.boxes = None
        self.trk_ids = None
        self.trk_pts = None
        self.line_thickness = 2
        self.trk_history = defaultdict(list)

        # Speed estimator information
        self.current_time = 0
        self.dist_data = {}
        self.trk_idslist = []
        self.spdl_dist_thresh = 10
        self.trk_previous_times = {}
        self.trk_previous_points = {}

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

    def set_args(
        self,
        reg_pts,
        names,
        view_img=False,
        line_thickness=2,
        region_thickness=5,
        spdl_dist_thresh=10,
    ):
        """
        Configures the speed estimation and display parameters.

        Args:
            reg_pts (list): Initial list of points defining the speed calculation region.
            names (dict): object detection classes names
            view_img (bool): Flag indicating frame display
            line_thickness (int): Line thickness for bounding boxes.
            region_thickness (int): Speed estimation region thickness
            spdl_dist_thresh (int): Euclidean distance threshold for speed line
        """
        if reg_pts is None:
            print("Region points not provided, using default values")
        else:
            self.reg_pts = reg_pts
        self.names = names
        self.view_img = view_img
        self.line_thickness = line_thickness
        self.region_thickness = region_thickness
        self.spdl_dist_thresh = spdl_dist_thresh

    def extract_tracks(self, tracks):
        """
        Extracts results from the provided data.

        Args:
            tracks (list): List of tracks obtained from the object tracking process.
        """
        self.boxes = tracks[0].boxes.xyxy.cpu()
        self.clss = tracks[0].boxes.cls.cpu().tolist()
        self.trk_ids = tracks[0].boxes.id.int().cpu().tolist()

    def store_track_info(self, track_id, box):
        """
        Store track data.

        Args:
            track_id (int): object track id.
            box (list): object bounding box data
        """
        track = self.trk_history[track_id]
        bbox_center = (float((box[0] + box[2]) / 2), float((box[1] + box[3]) / 2))
        track.append(bbox_center)

        if len(track) > 30:
            track.pop(0)

        self.trk_pts = np.hstack(track).astype(np.int32).reshape((-1, 1, 2))
        return track

    def plot_box_and_track(self, track_id, box, cls, track):
        """
        Plot track and bounding box.

        Args:
            track_id (int): object track id.
            box (list): object bounding box data
            cls (str): object class name
            track (list): tracking history for tracks path drawing
        """
        speed_label = f"{int(self.dist_data[track_id])}km/ph" if track_id in self.dist_data else self.names[int(cls)]
        bbox_color = colors(int(track_id)) if track_id in self.dist_data else (255, 0, 255)

        self.annotator.box_label(box, speed_label, bbox_color)

        cv2.polylines(self.im0, [self.trk_pts], isClosed=False, color=(0, 255, 0), thickness=1)
        cv2.circle(self.im0, (int(track[-1][0]), int(track[-1][1])), 5, bbox_color, -1)

    def calculate_speed(self, trk_id, track):
        """
        Calculation of object speed.

        Args:
            trk_id (int): object track id.
            track (list): tracking history for tracks path drawing
        """

        if not self.reg_pts[0][0] < track[-1][0] < self.reg_pts[1][0]:
            return
        if self.reg_pts[1][1] - self.spdl_dist_thresh < track[-1][1] < self.reg_pts[1][1] + self.spdl_dist_thresh:
            direction = "known"

        elif self.reg_pts[0][1] - self.spdl_dist_thresh < track[-1][1] < self.reg_pts[0][1] + self.spdl_dist_thresh:
            direction = "known"

        else:
            direction = "unknown"

        if self.trk_previous_times[trk_id] != 0 and direction != "unknown" and trk_id not in self.trk_idslist:
            self.trk_idslist.append(trk_id)

            time_difference = time() - self.trk_previous_times[trk_id]
            if time_difference > 0:
                dist_difference = np.abs(track[-1][1] - self.trk_previous_points[trk_id][1])
                speed = dist_difference / time_difference
                self.dist_data[trk_id] = speed

        self.trk_previous_times[trk_id] = time()
        self.trk_previous_points[trk_id] = track[-1]

    def estimate_speed(self, im0, tracks, region_color=(255, 0, 0)):
        """
        Calculate object based on tracking data.

        Args:
            im0 (nd array): Image
            tracks (list): List of tracks obtained from the object tracking process.
            region_color (tuple): Color to use when drawing regions.
        """
        self.im0 = im0
        if tracks[0].boxes.id is None:
            if self.view_img and self.env_check:
                self.display_frames()
            return im0
        self.extract_tracks(tracks)

        self.annotator = Annotator(self.im0, line_width=2)
        self.annotator.draw_region(reg_pts=self.reg_pts, color=region_color, thickness=self.region_thickness)

        for box, trk_id, cls in zip(self.boxes, self.trk_ids, self.clss):
            track = self.store_track_info(trk_id, box)

            if trk_id not in self.trk_previous_times:
                self.trk_previous_times[trk_id] = 0

            self.plot_box_and_track(trk_id, box, cls, track)
            self.calculate_speed(trk_id, track)

        if self.view_img and self.env_check:
            self.display_frames()

        return im0

    def display_frames(self):
        """Display frame."""
        cv2.imshow("Ultralytics Speed Estimation", self.im0)
        if cv2.waitKey(1) & 0xFF == ord("q"):
            return

__init__()

使用视觉、图像、轨迹和速度参数的默认值初始化速度估算器类。

源代码 ultralytics/solutions/speed_estimation.py
def __init__(self):
    """Initializes the speed-estimator class with default values for Visual, Image, track and speed parameters."""

    # Visual & im0 information
    self.im0 = None
    self.annotator = None
    self.view_img = False

    # Region information
    self.reg_pts = [(20, 400), (1260, 400)]
    self.region_thickness = 3

    # Predict/track information
    self.clss = None
    self.names = None
    self.boxes = None
    self.trk_ids = None
    self.trk_pts = None
    self.line_thickness = 2
    self.trk_history = defaultdict(list)

    # Speed estimator information
    self.current_time = 0
    self.dist_data = {}
    self.trk_idslist = []
    self.spdl_dist_thresh = 10
    self.trk_previous_times = {}
    self.trk_previous_points = {}

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

calculate_speed(trk_id, track)

计算物体速度

参数

名称 类型 说明 默认值
trk_id int

对象轨道 ID。

所需
track list

轨迹路径绘图的跟踪历史

所需
源代码 ultralytics/solutions/speed_estimation.py
def calculate_speed(self, trk_id, track):
    """
    Calculation of object speed.

    Args:
        trk_id (int): object track id.
        track (list): tracking history for tracks path drawing
    """

    if not self.reg_pts[0][0] < track[-1][0] < self.reg_pts[1][0]:
        return
    if self.reg_pts[1][1] - self.spdl_dist_thresh < track[-1][1] < self.reg_pts[1][1] + self.spdl_dist_thresh:
        direction = "known"

    elif self.reg_pts[0][1] - self.spdl_dist_thresh < track[-1][1] < self.reg_pts[0][1] + self.spdl_dist_thresh:
        direction = "known"

    else:
        direction = "unknown"

    if self.trk_previous_times[trk_id] != 0 and direction != "unknown" and trk_id not in self.trk_idslist:
        self.trk_idslist.append(trk_id)

        time_difference = time() - self.trk_previous_times[trk_id]
        if time_difference > 0:
            dist_difference = np.abs(track[-1][1] - self.trk_previous_points[trk_id][1])
            speed = dist_difference / time_difference
            self.dist_data[trk_id] = speed

    self.trk_previous_times[trk_id] = time()
    self.trk_previous_points[trk_id] = track[-1]

display_frames()

显示框。

源代码 ultralytics/solutions/speed_estimation.py
def display_frames(self):
    """Display frame."""
    cv2.imshow("Ultralytics Speed Estimation", self.im0)
    if cv2.waitKey(1) & 0xFF == ord("q"):
        return

estimate_speed(im0, tracks, region_color=(255, 0, 0))

根据跟踪数据计算目标。

参数

名称 类型 说明 默认值
im0 nd array

图片

所需
tracks list

物体追踪过程中获得的轨迹列表。

所需
region_color tuple

绘制区域时使用的颜色。

(255, 0, 0)
源代码 ultralytics/solutions/speed_estimation.py
def estimate_speed(self, im0, tracks, region_color=(255, 0, 0)):
    """
    Calculate object based on tracking data.

    Args:
        im0 (nd array): Image
        tracks (list): List of tracks obtained from the object tracking process.
        region_color (tuple): Color to use when drawing regions.
    """
    self.im0 = im0
    if tracks[0].boxes.id is None:
        if self.view_img and self.env_check:
            self.display_frames()
        return im0
    self.extract_tracks(tracks)

    self.annotator = Annotator(self.im0, line_width=2)
    self.annotator.draw_region(reg_pts=self.reg_pts, color=region_color, thickness=self.region_thickness)

    for box, trk_id, cls in zip(self.boxes, self.trk_ids, self.clss):
        track = self.store_track_info(trk_id, box)

        if trk_id not in self.trk_previous_times:
            self.trk_previous_times[trk_id] = 0

        self.plot_box_and_track(trk_id, box, cls, track)
        self.calculate_speed(trk_id, track)

    if self.view_img and self.env_check:
        self.display_frames()

    return im0

extract_tracks(tracks)

从提供的数据中提取结果。

参数

名称 类型 说明 默认值
tracks list

物体追踪过程中获得的轨迹列表。

所需
源代码 ultralytics/solutions/speed_estimation.py
def extract_tracks(self, tracks):
    """
    Extracts results from the provided data.

    Args:
        tracks (list): List of tracks obtained from the object tracking process.
    """
    self.boxes = tracks[0].boxes.xyxy.cpu()
    self.clss = tracks[0].boxes.cls.cpu().tolist()
    self.trk_ids = tracks[0].boxes.id.int().cpu().tolist()

plot_box_and_track(track_id, box, cls, track)

绘制轨迹和边界框。

参数

名称 类型 说明 默认值
track_id int

对象轨道 ID。

所需
box list

对象边界框数据

所需
cls str

对象类名

所需
track list

轨迹路径绘图的跟踪历史

所需
源代码 ultralytics/solutions/speed_estimation.py
def plot_box_and_track(self, track_id, box, cls, track):
    """
    Plot track and bounding box.

    Args:
        track_id (int): object track id.
        box (list): object bounding box data
        cls (str): object class name
        track (list): tracking history for tracks path drawing
    """
    speed_label = f"{int(self.dist_data[track_id])}km/ph" if track_id in self.dist_data else self.names[int(cls)]
    bbox_color = colors(int(track_id)) if track_id in self.dist_data else (255, 0, 255)

    self.annotator.box_label(box, speed_label, bbox_color)

    cv2.polylines(self.im0, [self.trk_pts], isClosed=False, color=(0, 255, 0), thickness=1)
    cv2.circle(self.im0, (int(track[-1][0]), int(track[-1][1])), 5, bbox_color, -1)

set_args(reg_pts, names, view_img=False, line_thickness=2, region_thickness=5, spdl_dist_thresh=10)

配置速度估算和显示参数。

参数

名称 类型 说明 默认值
reg_pts list

定义速度计算区域的初始点列表。

所需
names dict

对象检测类名称

所需
view_img bool

显示帧的标志

False
line_thickness int

边界框的线条粗细

2
region_thickness int

速度估计区域厚度

5
spdl_dist_thresh int

速度线的欧氏距离阈值

10
源代码 ultralytics/solutions/speed_estimation.py
def set_args(
    self,
    reg_pts,
    names,
    view_img=False,
    line_thickness=2,
    region_thickness=5,
    spdl_dist_thresh=10,
):
    """
    Configures the speed estimation and display parameters.

    Args:
        reg_pts (list): Initial list of points defining the speed calculation region.
        names (dict): object detection classes names
        view_img (bool): Flag indicating frame display
        line_thickness (int): Line thickness for bounding boxes.
        region_thickness (int): Speed estimation region thickness
        spdl_dist_thresh (int): Euclidean distance threshold for speed line
    """
    if reg_pts is None:
        print("Region points not provided, using default values")
    else:
        self.reg_pts = reg_pts
    self.names = names
    self.view_img = view_img
    self.line_thickness = line_thickness
    self.region_thickness = region_thickness
    self.spdl_dist_thresh = spdl_dist_thresh

store_track_info(track_id, box)

存储轨道数据

参数

名称 类型 说明 默认值
track_id int

对象轨道 ID。

所需
box list

对象边界框数据

所需
源代码 ultralytics/solutions/speed_estimation.py
def store_track_info(self, track_id, box):
    """
    Store track data.

    Args:
        track_id (int): object track id.
        box (list): object bounding box data
    """
    track = self.trk_history[track_id]
    bbox_center = (float((box[0] + box[2]) / 2), float((box[1] + box[3]) / 2))
    track.append(bbox_center)

    if len(track) > 30:
        track.pop(0)

    self.trk_pts = np.hstack(track).astype(np.int32).reshape((-1, 1, 2))
    return track





创建于 2024-01-05,更新于 2024-01-10
作者:AyushExel(1),chr043416@gmail.com(1)