انتقل إلى المحتوى

مرجع ل ultralytics/solutions/distance_calculation.py

ملاحظه

هذا الملف متاح في https://github.com/ultralytics/ultralytics/ نقطة / الرئيسية /ultralytics/الحلول/distance_calculation.py. إذا اكتشفت مشكلة ، فيرجى المساعدة في إصلاحها من خلال المساهمة في طلب 🛠️ سحب. شكرا لك 🙏!



ultralytics.solutions.distance_calculation.DistanceCalculation

فئة لحساب المسافة بين كائنين في دفق الفيديو في الوقت الفعلي بناء على مساراتهما.

شفرة المصدر في ultralytics/solutions/distance_calculation.py
class DistanceCalculation:
    """A class to calculate distance between two objects in real-time video stream based on their tracks."""

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

        # Visual & im0 information
        self.im0 = None
        self.annotator = None
        self.view_img = False
        self.line_color = (255, 255, 0)
        self.centroid_color = (255, 0, 255)

        # Predict/track information
        self.clss = None
        self.names = None
        self.boxes = None
        self.line_thickness = 2
        self.trk_ids = None

        # Distance calculation information
        self.centroids = []
        self.pixel_per_meter = 10

        # Mouse event
        self.left_mouse_count = 0
        self.selected_boxes = {}

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

    def set_args(
        self,
        names,
        pixels_per_meter=10,
        view_img=False,
        line_thickness=2,
        line_color=(255, 255, 0),
        centroid_color=(255, 0, 255),
    ):
        """
        Configures the distance calculation and display parameters.

        Args:
            names (dict): object detection classes names
            pixels_per_meter (int): Number of pixels in meter
            view_img (bool): Flag indicating frame display
            line_thickness (int): Line thickness for bounding boxes.
            line_color (RGB): color of centroids line
            centroid_color (RGB): colors of bbox centroids
        """
        self.names = names
        self.pixel_per_meter = pixels_per_meter
        self.view_img = view_img
        self.line_thickness = line_thickness
        self.line_color = line_color
        self.centroid_color = centroid_color

    def mouse_event_for_distance(self, event, x, y, flags, param):
        """
        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.).
            param (dict): Additional parameters you may want to pass to the function.
        """
        global selected_boxes
        global left_mouse_count
        if event == cv2.EVENT_LBUTTONDOWN:
            self.left_mouse_count += 1
            if self.left_mouse_count <= 2:
                for box, track_id in zip(self.boxes, self.trk_ids):
                    if box[0] < x < box[2] and box[1] < y < box[3] and track_id not in self.selected_boxes:
                        self.selected_boxes[track_id] = []
                        self.selected_boxes[track_id] = box

        if event == cv2.EVENT_RBUTTONDOWN:
            self.selected_boxes = {}
            self.left_mouse_count = 0

    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 calculate_centroid(self, box):
        """
        Calculate the centroid of bounding box.

        Args:
            box (list): Bounding box data
        """
        return int((box[0] + box[2]) // 2), int((box[1] + box[3]) // 2)

    def calculate_distance(self, centroid1, centroid2):
        """
        Calculate distance between two centroids.

        Args:
            centroid1 (point): First bounding box data
            centroid2 (point): Second bounding box data
        """
        pixel_distance = math.sqrt((centroid1[0] - centroid2[0]) ** 2 + (centroid1[1] - centroid2[1]) ** 2)
        return pixel_distance / self.pixel_per_meter, (pixel_distance / self.pixel_per_meter) * 1000

    def start_process(self, im0, tracks):
        """
        Calculate distance between two bounding boxes based on tracking data.

        Args:
            im0 (nd array): Image
            tracks (list): List of tracks obtained from the object tracking process.
        """
        self.im0 = im0
        if tracks[0].boxes.id is None:
            if self.view_img:
                self.display_frames()
            return
        self.extract_tracks(tracks)

        self.annotator = Annotator(self.im0, line_width=2)

        for box, cls, track_id in zip(self.boxes, self.clss, self.trk_ids):
            self.annotator.box_label(box, color=colors(int(cls), True), label=self.names[int(cls)])

            if len(self.selected_boxes) == 2:
                for trk_id, _ in self.selected_boxes.items():
                    if trk_id == track_id:
                        self.selected_boxes[track_id] = box

        if len(self.selected_boxes) == 2:
            for trk_id, box in self.selected_boxes.items():
                centroid = self.calculate_centroid(self.selected_boxes[trk_id])
                self.centroids.append(centroid)

            distance_m, distance_mm = self.calculate_distance(self.centroids[0], self.centroids[1])
            self.annotator.plot_distance_and_line(
                distance_m, distance_mm, self.centroids, self.line_color, self.centroid_color
            )

        self.centroids = []

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

        return im0

    def display_frames(self):
        """Display frame."""
        cv2.namedWindow("Ultralytics Distance Estimation")
        cv2.setMouseCallback("Ultralytics Distance Estimation", self.mouse_event_for_distance)
        cv2.imshow("Ultralytics Distance Estimation", self.im0)

        if cv2.waitKey(1) & 0xFF == ord("q"):
            return

__init__()

تهيئة فئة حساب المسافة بالقيم الافتراضية للمرئيات والصورة والمسار والمسافة البارامترات.

شفرة المصدر في ultralytics/solutions/distance_calculation.py
14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 2930 31 32 3334 35 36 373839404142
def __init__(self):
    """Initializes the distance calculation class with default values for Visual, Image, track and distance
    parameters.
    """

    # Visual & im0 information
    self.im0 = None
    self.annotator = None
    self.view_img = False
    self.line_color = (255, 255, 0)
    self.centroid_color = (255, 0, 255)

    # Predict/track information
    self.clss = None
    self.names = None
    self.boxes = None
    self.line_thickness = 2
    self.trk_ids = None

    # Distance calculation information
    self.centroids = []
    self.pixel_per_meter = 10

    # Mouse event
    self.left_mouse_count = 0
    self.selected_boxes = {}

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

calculate_centroid(box)

احسب النقطه الوسطى للمربع المحيط.

البارامترات:

اسم نوع وصف افتراضي
box list

بيانات المربع المحيط

مطلوب
شفرة المصدر في ultralytics/solutions/distance_calculation.py
108 109 110 111 112 113 114115
def calculate_centroid(self, box):
    """
    Calculate the centroid of bounding box.

    Args:
        box (list): Bounding box data
    """
    return int((box[0] + box[2]) // 2), int((box[1] + box[3]) // 2)

calculate_distance(centroid1, centroid2)

احسب المسافة بين اثنين من السنترويدات.

البارامترات:

اسم نوع وصف افتراضي
centroid1 point

بيانات المربع المحيط الأول

مطلوب
centroid2 point

بيانات المربع المحيط الثاني

مطلوب
شفرة المصدر في ultralytics/solutions/distance_calculation.py
117118 119 120 121 122 123 124 125126
def calculate_distance(self, centroid1, centroid2):
    """
    Calculate distance between two centroids.

    Args:
        centroid1 (point): First bounding box data
        centroid2 (point): Second bounding box data
    """
    pixel_distance = math.sqrt((centroid1[0] - centroid2[0]) ** 2 + (centroid1[1] - centroid2[1]) ** 2)
    return pixel_distance / self.pixel_per_meter, (pixel_distance / self.pixel_per_meter) * 1000

display_frames()

إطار العرض.

شفرة المصدر في ultralytics/solutions/distance_calculation.py
170 171 172 173 174 175 176 177
def display_frames(self):
    """Display frame."""
    cv2.namedWindow("Ultralytics Distance Estimation")
    cv2.setMouseCallback("Ultralytics Distance Estimation", self.mouse_event_for_distance)
    cv2.imshow("Ultralytics Distance Estimation", self.im0)

    if cv2.waitKey(1) & 0xFF == ord("q"):
        return

extract_tracks(tracks)

استخراج النتائج من البيانات المقدمة.

البارامترات:

اسم نوع وصف افتراضي
tracks list

قائمة المسارات التي تم الحصول عليها من عملية تتبع الكائن.

مطلوب
شفرة المصدر في ultralytics/solutions/distance_calculation.py
 97 98 99 100 101 102 103 104 105 106
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()

mouse_event_for_distance(event, x, y, flags, param)

تم تصميم هذه الوظيفة لنقل المنطقة مع أحداث الماوس في دفق فيديو في الوقت الفعلي.

البارامترات:

اسم نوع وصف افتراضي
event int

نوع حدث الماوس (على سبيل المثال، cv2. EVENT_MOUSEMOVE ، السيرة الذاتية 2. EVENT_LBUTTONDOWN ، إلخ).

مطلوب
x int

إحداثيات x لمؤشر الماوس.

مطلوب
y int

إحداثيات y لمؤشر الماوس.

مطلوب
flags int

أي علامات مرتبطة بالحدث (على سبيل المثال، cv2. EVENT_FLAG_CTRLKEY، السيرة الذاتية 2. EVENT_FLAG_SHIFTKEY ، وما إلى ذلك).

مطلوب
param dict

معلمات إضافية قد ترغب في تمريرها إلى الوظيفة.

مطلوب
شفرة المصدر في ultralytics/solutions/distance_calculation.py
71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 8687888990 9192939495
def mouse_event_for_distance(self, event, x, y, flags, param):
    """
    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.).
        param (dict): Additional parameters you may want to pass to the function.
    """
    global selected_boxes
    global left_mouse_count
    if event == cv2.EVENT_LBUTTONDOWN:
        self.left_mouse_count += 1
        if self.left_mouse_count <= 2:
            for box, track_id in zip(self.boxes, self.trk_ids):
                if box[0] < x < box[2] and box[1] < y < box[3] and track_id not in self.selected_boxes:
                    self.selected_boxes[track_id] = []
                    self.selected_boxes[track_id] = box

    if event == cv2.EVENT_RBUTTONDOWN:
        self.selected_boxes = {}
        self.left_mouse_count = 0

set_args(names, pixels_per_meter=10, view_img=False, line_thickness=2, line_color=(255, 255, 0), centroid_color=(255, 0, 255))

يقوم بتكوين حساب المسافة ومعلمات العرض.

البارامترات:

اسم نوع وصف افتراضي
names dict

أسماء فئات الكشف عن الكائنات

مطلوب
pixels_per_meter int

عدد البكسل بالمتر

10
view_img bool

علامة تشير إلى عرض الإطار

False
line_thickness int

سمك الخط للمربعات المحيطة.

2
line_color RGB

لون خط السنترويدات

(255, 255, 0)
centroid_color RGB

ألوان bbox centroids

(255, 0, 255)
شفرة المصدر في ultralytics/solutions/distance_calculation.py
44 45 46 47 48 49 50 51 52 53 54 55 56 5758596061626364 6566 67 68 69
def set_args(
    self,
    names,
    pixels_per_meter=10,
    view_img=False,
    line_thickness=2,
    line_color=(255, 255, 0),
    centroid_color=(255, 0, 255),
):
    """
    Configures the distance calculation and display parameters.

    Args:
        names (dict): object detection classes names
        pixels_per_meter (int): Number of pixels in meter
        view_img (bool): Flag indicating frame display
        line_thickness (int): Line thickness for bounding boxes.
        line_color (RGB): color of centroids line
        centroid_color (RGB): colors of bbox centroids
    """
    self.names = names
    self.pixel_per_meter = pixels_per_meter
    self.view_img = view_img
    self.line_thickness = line_thickness
    self.line_color = line_color
    self.centroid_color = centroid_color

start_process(im0, tracks)

احسب المسافة بين مربعين محيطين بناء على بيانات التتبع.

البارامترات:

اسم نوع وصف افتراضي
im0 nd array

صورة

مطلوب
tracks list

قائمة المسارات التي تم الحصول عليها من عملية تتبع الكائن.

مطلوب
شفرة المصدر في ultralytics/solutions/distance_calculation.py
128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158159160 161 162 163 164 165 166 167 168
def start_process(self, im0, tracks):
    """
    Calculate distance between two bounding boxes based on tracking data.

    Args:
        im0 (nd array): Image
        tracks (list): List of tracks obtained from the object tracking process.
    """
    self.im0 = im0
    if tracks[0].boxes.id is None:
        if self.view_img:
            self.display_frames()
        return
    self.extract_tracks(tracks)

    self.annotator = Annotator(self.im0, line_width=2)

    for box, cls, track_id in zip(self.boxes, self.clss, self.trk_ids):
        self.annotator.box_label(box, color=colors(int(cls), True), label=self.names[int(cls)])

        if len(self.selected_boxes) == 2:
            for trk_id, _ in self.selected_boxes.items():
                if trk_id == track_id:
                    self.selected_boxes[track_id] = box

    if len(self.selected_boxes) == 2:
        for trk_id, box in self.selected_boxes.items():
            centroid = self.calculate_centroid(self.selected_boxes[trk_id])
            self.centroids.append(centroid)

        distance_m, distance_mm = self.calculate_distance(self.centroids[0], self.centroids[1])
        self.annotator.plot_distance_and_line(
            distance_m, distance_mm, self.centroids, self.line_color, self.centroid_color
        )

    self.centroids = []

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

    return im0





تم إنشاؤه في 2024-01-05, اخر تحديث 2024-01-10
المؤلفون: أيوشكسل (1) ، chr043416@gmail.com (1)