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

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

ملاحظه

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



ultralytics.solutions.heatmap.Heatmap

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

شفرة المصدر في ultralytics/solutions/heatmap.py
class Heatmap:
    """A class to draw heatmaps in real-time video stream based on their tracks."""

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

        # Visual information
        self.annotator = None
        self.view_img = False
        self.shape = "circle"

        # Image information
        self.imw = None
        self.imh = None
        self.im0 = None
        self.view_in_counts = True
        self.view_out_counts = True

        # Heatmap colormap and heatmap np array
        self.colormap = None
        self.heatmap = None
        self.heatmap_alpha = 0.5

        # Predict/track information
        self.boxes = None
        self.track_ids = None
        self.clss = None
        self.track_history = defaultdict(list)

        # Region & Line Information
        self.count_reg_pts = None
        self.counting_region = None
        self.line_dist_thresh = 15
        self.region_thickness = 5
        self.region_color = (255, 0, 255)

        # Object Counting Information
        self.in_counts = 0
        self.out_counts = 0
        self.counting_list = []
        self.count_txt_thickness = 0
        self.count_txt_color = (0, 0, 0)
        self.count_color = (255, 255, 255)

        # Decay factor
        self.decay_factor = 0.99

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

    def set_args(
        self,
        imw,
        imh,
        colormap=cv2.COLORMAP_JET,
        heatmap_alpha=0.5,
        view_img=False,
        view_in_counts=True,
        view_out_counts=True,
        count_reg_pts=None,
        count_txt_thickness=2,
        count_txt_color=(0, 0, 0),
        count_color=(255, 255, 255),
        count_reg_color=(255, 0, 255),
        region_thickness=5,
        line_dist_thresh=15,
        decay_factor=0.99,
        shape="circle",
    ):
        """
        Configures the heatmap colormap, width, height and display parameters.

        Args:
            colormap (cv2.COLORMAP): The colormap to be set.
            imw (int): The width of the frame.
            imh (int): The height of the frame.
            heatmap_alpha (float): alpha value for heatmap display
            view_img (bool): Flag indicating frame display
            view_in_counts (bool): Flag to control whether to display the incounts on video stream.
            view_out_counts (bool): Flag to control whether to display the outcounts on video stream.
            count_reg_pts (list): Object counting region points
            count_txt_thickness (int): Text thickness for object counting display
            count_txt_color (RGB color): count text color value
            count_color (RGB color): count text background color value
            count_reg_color (RGB color): Color of object counting region
            region_thickness (int): Object counting Region thickness
            line_dist_thresh (int): Euclidean Distance threshold for line counter
            decay_factor (float): value for removing heatmap area after object passed
            shape (str): Heatmap shape, rect or circle shape supported
        """
        self.imw = imw
        self.imh = imh
        self.heatmap_alpha = heatmap_alpha
        self.view_img = view_img
        self.view_in_counts = view_in_counts
        self.view_out_counts = view_out_counts
        self.colormap = colormap

        # Region and line selection
        if count_reg_pts is not None:
            if len(count_reg_pts) == 2:
                print("Line Counter Initiated.")
                self.count_reg_pts = count_reg_pts
                self.counting_region = LineString(count_reg_pts)

            elif len(count_reg_pts) == 4:
                print("Region Counter Initiated.")
                self.count_reg_pts = count_reg_pts
                self.counting_region = Polygon(self.count_reg_pts)

            else:
                print("Region or line points Invalid, 2 or 4 points supported")
                print("Using Line Counter Now")
                self.counting_region = Polygon([(20, 400), (1260, 400)])  # dummy points

        # Heatmap new frame
        self.heatmap = np.zeros((int(self.imh), int(self.imw)), dtype=np.float32)

        self.count_txt_thickness = count_txt_thickness
        self.count_txt_color = count_txt_color
        self.count_color = count_color
        self.region_color = count_reg_color
        self.region_thickness = region_thickness
        self.decay_factor = decay_factor
        self.line_dist_thresh = line_dist_thresh
        self.shape = shape

        # shape of heatmap, if not selected
        if self.shape not in ["circle", "rect"]:
            print("Unknown shape value provided, 'circle' & 'rect' supported")
            print("Using Circular shape now")
            self.shape = "circle"

    def extract_results(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.track_ids = tracks[0].boxes.id.int().cpu().tolist()

    def generate_heatmap(self, im0, tracks):
        """
        Generate heatmap 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:
            self.heatmap = np.zeros((int(self.imh), int(self.imw)), dtype=np.float32)
            if self.view_img and self.env_check:
                self.display_frames()
            return im0
        self.heatmap *= self.decay_factor  # decay factor
        self.extract_results(tracks)
        self.annotator = Annotator(self.im0, self.count_txt_thickness, None)

        if self.count_reg_pts is not None:
            # Draw counting region
            if self.view_in_counts or self.view_out_counts:
                self.annotator.draw_region(
                    reg_pts=self.count_reg_pts, color=self.region_color, thickness=self.region_thickness
                )

            for box, cls, track_id in zip(self.boxes, self.clss, self.track_ids):
                if self.shape == "circle":
                    center = (int((box[0] + box[2]) // 2), int((box[1] + box[3]) // 2))
                    radius = min(int(box[2]) - int(box[0]), int(box[3]) - int(box[1])) // 2

                    y, x = np.ogrid[0 : self.heatmap.shape[0], 0 : self.heatmap.shape[1]]
                    mask = (x - center[0]) ** 2 + (y - center[1]) ** 2 <= radius**2

                    self.heatmap[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])] += (
                        2 * mask[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])]
                    )

                else:
                    self.heatmap[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])] += 2

                # Store tracking hist
                track_line = self.track_history[track_id]
                track_line.append((float((box[0] + box[2]) / 2), float((box[1] + box[3]) / 2)))
                if len(track_line) > 30:
                    track_line.pop(0)

                # Count objects
                if len(self.count_reg_pts) == 4:
                    if self.counting_region.contains(Point(track_line[-1])) and track_id not in self.counting_list:
                        self.counting_list.append(track_id)
                        if box[0] < self.counting_region.centroid.x:
                            self.out_counts += 1
                        else:
                            self.in_counts += 1

                elif len(self.count_reg_pts) == 2:
                    distance = Point(track_line[-1]).distance(self.counting_region)
                    if distance < self.line_dist_thresh and track_id not in self.counting_list:
                        self.counting_list.append(track_id)
                        if box[0] < self.counting_region.centroid.x:
                            self.out_counts += 1
                        else:
                            self.in_counts += 1
        else:
            for box, cls in zip(self.boxes, self.clss):
                if self.shape == "circle":
                    center = (int((box[0] + box[2]) // 2), int((box[1] + box[3]) // 2))
                    radius = min(int(box[2]) - int(box[0]), int(box[3]) - int(box[1])) // 2

                    y, x = np.ogrid[0 : self.heatmap.shape[0], 0 : self.heatmap.shape[1]]
                    mask = (x - center[0]) ** 2 + (y - center[1]) ** 2 <= radius**2

                    self.heatmap[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])] += (
                        2 * mask[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])]
                    )

                else:
                    self.heatmap[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])] += 2

        # Normalize, apply colormap to heatmap and combine with original image
        heatmap_normalized = cv2.normalize(self.heatmap, None, 0, 255, cv2.NORM_MINMAX)
        heatmap_colored = cv2.applyColorMap(heatmap_normalized.astype(np.uint8), self.colormap)

        incount_label = f"In Count : {self.in_counts}"
        outcount_label = f"OutCount : {self.out_counts}"

        # Display counts based on user choice
        counts_label = None
        if not self.view_in_counts and not self.view_out_counts:
            counts_label = None
        elif not self.view_in_counts:
            counts_label = outcount_label
        elif not self.view_out_counts:
            counts_label = incount_label
        else:
            counts_label = f"{incount_label} {outcount_label}"

        if self.count_reg_pts is not None and counts_label is not None:
            self.annotator.count_labels(
                counts=counts_label,
                count_txt_size=self.count_txt_thickness,
                txt_color=self.count_txt_color,
                color=self.count_color,
            )

        self.im0 = cv2.addWeighted(self.im0, 1 - self.heatmap_alpha, heatmap_colored, self.heatmap_alpha, 0)

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

        return self.im0

    def display_frames(self):
        """Display frame."""
        cv2.imshow("Ultralytics Heatmap", self.im0)

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

__init__()

تهيئة فئة خريطة الحرارة بالقيم الافتراضية لمعلمات الصورة والصورة والتعقب والعد وخريطة الحرارة.

شفرة المصدر في ultralytics/solutions/heatmap.py
19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 4344454647484950515253 54 555657585960616263 64
def __init__(self):
    """Initializes the heatmap class with default values for Visual, Image, track, count and heatmap parameters."""

    # Visual information
    self.annotator = None
    self.view_img = False
    self.shape = "circle"

    # Image information
    self.imw = None
    self.imh = None
    self.im0 = None
    self.view_in_counts = True
    self.view_out_counts = True

    # Heatmap colormap and heatmap np array
    self.colormap = None
    self.heatmap = None
    self.heatmap_alpha = 0.5

    # Predict/track information
    self.boxes = None
    self.track_ids = None
    self.clss = None
    self.track_history = defaultdict(list)

    # Region & Line Information
    self.count_reg_pts = None
    self.counting_region = None
    self.line_dist_thresh = 15
    self.region_thickness = 5
    self.region_color = (255, 0, 255)

    # Object Counting Information
    self.in_counts = 0
    self.out_counts = 0
    self.counting_list = []
    self.count_txt_thickness = 0
    self.count_txt_color = (0, 0, 0)
    self.count_color = (255, 255, 255)

    # Decay factor
    self.decay_factor = 0.99

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

display_frames()

إطار العرض.

شفرة المصدر في ultralytics/solutions/heatmap.py
272 273 274 275 276 277
def display_frames(self):
    """Display frame."""
    cv2.imshow("Ultralytics Heatmap", self.im0)

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

extract_results(tracks)

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

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

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

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

مطلوب
شفرة المصدر في ultralytics/solutions/heatmap.py
149 150 151 152 153 154 155 156 157 158
def extract_results(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.track_ids = tracks[0].boxes.id.int().cpu().tolist()

generate_heatmap(im0, tracks)

إنشاء خريطة حرارية بناء على بيانات التتبع.

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

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

صورة

مطلوب
tracks list

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

مطلوب
شفرة المصدر في ultralytics/solutions/heatmap.py
160 161 162 163 164 165 166 167 168169 170 171 172 173 174 175 176 177 178179 180 181 182 183 184 185 186 187188189 190191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208209 210 211 212 213 214 215 216 217218 219 220 221222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258259 260 261 262 263 264 265 266 267 268 269270
def generate_heatmap(self, im0, tracks):
    """
    Generate heatmap 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:
        self.heatmap = np.zeros((int(self.imh), int(self.imw)), dtype=np.float32)
        if self.view_img and self.env_check:
            self.display_frames()
        return im0
    self.heatmap *= self.decay_factor  # decay factor
    self.extract_results(tracks)
    self.annotator = Annotator(self.im0, self.count_txt_thickness, None)

    if self.count_reg_pts is not None:
        # Draw counting region
        if self.view_in_counts or self.view_out_counts:
            self.annotator.draw_region(
                reg_pts=self.count_reg_pts, color=self.region_color, thickness=self.region_thickness
            )

        for box, cls, track_id in zip(self.boxes, self.clss, self.track_ids):
            if self.shape == "circle":
                center = (int((box[0] + box[2]) // 2), int((box[1] + box[3]) // 2))
                radius = min(int(box[2]) - int(box[0]), int(box[3]) - int(box[1])) // 2

                y, x = np.ogrid[0 : self.heatmap.shape[0], 0 : self.heatmap.shape[1]]
                mask = (x - center[0]) ** 2 + (y - center[1]) ** 2 <= radius**2

                self.heatmap[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])] += (
                    2 * mask[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])]
                )

            else:
                self.heatmap[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])] += 2

            # Store tracking hist
            track_line = self.track_history[track_id]
            track_line.append((float((box[0] + box[2]) / 2), float((box[1] + box[3]) / 2)))
            if len(track_line) > 30:
                track_line.pop(0)

            # Count objects
            if len(self.count_reg_pts) == 4:
                if self.counting_region.contains(Point(track_line[-1])) and track_id not in self.counting_list:
                    self.counting_list.append(track_id)
                    if box[0] < self.counting_region.centroid.x:
                        self.out_counts += 1
                    else:
                        self.in_counts += 1

            elif len(self.count_reg_pts) == 2:
                distance = Point(track_line[-1]).distance(self.counting_region)
                if distance < self.line_dist_thresh and track_id not in self.counting_list:
                    self.counting_list.append(track_id)
                    if box[0] < self.counting_region.centroid.x:
                        self.out_counts += 1
                    else:
                        self.in_counts += 1
    else:
        for box, cls in zip(self.boxes, self.clss):
            if self.shape == "circle":
                center = (int((box[0] + box[2]) // 2), int((box[1] + box[3]) // 2))
                radius = min(int(box[2]) - int(box[0]), int(box[3]) - int(box[1])) // 2

                y, x = np.ogrid[0 : self.heatmap.shape[0], 0 : self.heatmap.shape[1]]
                mask = (x - center[0]) ** 2 + (y - center[1]) ** 2 <= radius**2

                self.heatmap[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])] += (
                    2 * mask[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])]
                )

            else:
                self.heatmap[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])] += 2

    # Normalize, apply colormap to heatmap and combine with original image
    heatmap_normalized = cv2.normalize(self.heatmap, None, 0, 255, cv2.NORM_MINMAX)
    heatmap_colored = cv2.applyColorMap(heatmap_normalized.astype(np.uint8), self.colormap)

    incount_label = f"In Count : {self.in_counts}"
    outcount_label = f"OutCount : {self.out_counts}"

    # Display counts based on user choice
    counts_label = None
    if not self.view_in_counts and not self.view_out_counts:
        counts_label = None
    elif not self.view_in_counts:
        counts_label = outcount_label
    elif not self.view_out_counts:
        counts_label = incount_label
    else:
        counts_label = f"{incount_label} {outcount_label}"

    if self.count_reg_pts is not None and counts_label is not None:
        self.annotator.count_labels(
            counts=counts_label,
            count_txt_size=self.count_txt_thickness,
            txt_color=self.count_txt_color,
            color=self.count_color,
        )

    self.im0 = cv2.addWeighted(self.im0, 1 - self.heatmap_alpha, heatmap_colored, self.heatmap_alpha, 0)

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

    return self.im0

set_args(imw, imh, colormap=cv2.COLORMAP_JET, heatmap_alpha=0.5, view_img=False, view_in_counts=True, view_out_counts=True, count_reg_pts=None, count_txt_thickness=2, count_txt_color=(0, 0, 0), count_color=(255, 255, 255), count_reg_color=(255, 0, 255), region_thickness=5, line_dist_thresh=15, decay_factor=0.99, shape='circle')

يقوم بتكوين خريطة ألوان خريطة التمثيل اللوني والعرض والارتفاع ومعلمات العرض.

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

اسم نوع وصف افتراضي
colormap COLORMAP

خريطة الألوان المراد تعيينها.

COLORMAP_JET
imw int

عرض الإطار.

مطلوب
imh int

ارتفاع الإطار.

مطلوب
heatmap_alpha float

قيمة ألفا لعرض خريطة التمثيل اللوني

0.5
view_img bool

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

False
view_in_counts bool

ضع علامة للتحكم في ما إذا كنت تريد عرض الأعداد في دفق الفيديو.

True
view_out_counts bool

ضع علامة للتحكم في ما إذا كنت تريد عرض الأعداد الخارجة في دفق الفيديو.

True
count_reg_pts list

نقاط منطقة عد الكائنات

None
count_txt_thickness int

سمك النص لعرض عد الكائنات

2
count_txt_color RGB color

عد قيمة لون النص

(0, 0, 0)
count_color RGB color

عد قيمة لون خلفية النص

(255, 255, 255)
count_reg_color RGB color

لون منطقة عد الكائنات

(255, 0, 255)
region_thickness int

عد الكائنات سمك المنطقة

5
line_dist_thresh int

عتبة المسافة الإقليدية لعداد الخط

15
decay_factor float

قيمة لإزالة منطقة خريطة التمثيل اللوني بعد تمرير الكائن

0.99
shape str

شكل خريطة التمثيل اللوني أو الشكل المستقيم أو الدائري مدعوم

'circle'
شفرة المصدر في ultralytics/solutions/heatmap.py
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 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130131 132 133 134 135 136 137 138139 140 141 142 143 144 145146 147
def set_args(
    self,
    imw,
    imh,
    colormap=cv2.COLORMAP_JET,
    heatmap_alpha=0.5,
    view_img=False,
    view_in_counts=True,
    view_out_counts=True,
    count_reg_pts=None,
    count_txt_thickness=2,
    count_txt_color=(0, 0, 0),
    count_color=(255, 255, 255),
    count_reg_color=(255, 0, 255),
    region_thickness=5,
    line_dist_thresh=15,
    decay_factor=0.99,
    shape="circle",
):
    """
    Configures the heatmap colormap, width, height and display parameters.

    Args:
        colormap (cv2.COLORMAP): The colormap to be set.
        imw (int): The width of the frame.
        imh (int): The height of the frame.
        heatmap_alpha (float): alpha value for heatmap display
        view_img (bool): Flag indicating frame display
        view_in_counts (bool): Flag to control whether to display the incounts on video stream.
        view_out_counts (bool): Flag to control whether to display the outcounts on video stream.
        count_reg_pts (list): Object counting region points
        count_txt_thickness (int): Text thickness for object counting display
        count_txt_color (RGB color): count text color value
        count_color (RGB color): count text background color value
        count_reg_color (RGB color): Color of object counting region
        region_thickness (int): Object counting Region thickness
        line_dist_thresh (int): Euclidean Distance threshold for line counter
        decay_factor (float): value for removing heatmap area after object passed
        shape (str): Heatmap shape, rect or circle shape supported
    """
    self.imw = imw
    self.imh = imh
    self.heatmap_alpha = heatmap_alpha
    self.view_img = view_img
    self.view_in_counts = view_in_counts
    self.view_out_counts = view_out_counts
    self.colormap = colormap

    # Region and line selection
    if count_reg_pts is not None:
        if len(count_reg_pts) == 2:
            print("Line Counter Initiated.")
            self.count_reg_pts = count_reg_pts
            self.counting_region = LineString(count_reg_pts)

        elif len(count_reg_pts) == 4:
            print("Region Counter Initiated.")
            self.count_reg_pts = count_reg_pts
            self.counting_region = Polygon(self.count_reg_pts)

        else:
            print("Region or line points Invalid, 2 or 4 points supported")
            print("Using Line Counter Now")
            self.counting_region = Polygon([(20, 400), (1260, 400)])  # dummy points

    # Heatmap new frame
    self.heatmap = np.zeros((int(self.imh), int(self.imw)), dtype=np.float32)

    self.count_txt_thickness = count_txt_thickness
    self.count_txt_color = count_txt_color
    self.count_color = count_color
    self.region_color = count_reg_color
    self.region_thickness = region_thickness
    self.decay_factor = decay_factor
    self.line_dist_thresh = line_dist_thresh
    self.shape = shape

    # shape of heatmap, if not selected
    if self.shape not in ["circle", "rect"]:
        print("Unknown shape value provided, 'circle' & 'rect' supported")
        print("Using Circular shape now")
        self.shape = "circle"





تم إنشاؤه في 2023-12-10, اخر تحديث 2023-12-10
المؤلفون: جلين جوشر (1)