ultralytics 8.1.29__py3-none-any.whl → 8.3.62__py3-none-any.whl

ultralytics/solutions/distance_calculation.py

@@ -1,181 +1,124 @@
-# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 import math
 
 import cv2
 
-from ultralytics.utils.checks import check_imshow
+from ultralytics.solutions.solutions import BaseSolution
 from ultralytics.utils.plotting import Annotator, colors
 
 
-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
+class DistanceCalculation(BaseSolution):
+    """
+    A class to calculate distance between two objects in a real-time video stream based on their tracks.
+
+    This class extends BaseSolution to provide functionality for selecting objects and calculating the distance
+    between them in a video stream using YOLO object detection and tracking.
+
+    Attributes:
+        left_mouse_count (int): Counter for left mouse button clicks.
+        selected_boxes (Dict[int, List[float]]): Dictionary to store selected bounding boxes and their track IDs.
+        annotator (Annotator): An instance of the Annotator class for drawing on the image.
+        boxes (List[List[float]]): List of bounding boxes for detected objects.
+        track_ids (List[int]): List of track IDs for detected objects.
+        clss (List[int]): List of class indices for detected objects.
+        names (List[str]): List of class names that the model can detect.
+        centroids (List[List[int]]): List to store centroids of selected bounding boxes.
+
+    Methods:
+        mouse_event_for_distance: Handles mouse events for selecting objects in the video stream.
+        calculate: Processes video frames and calculates the distance between selected objects.
+
+    Examples:
+        >>> distance_calc = DistanceCalculation()
+        >>> frame = cv2.imread("frame.jpg")
+        >>> processed_frame = distance_calc.calculate(frame)
+        >>> cv2.imshow("Distance Calculation", processed_frame)
+        >>> cv2.waitKey(0)
+    """
+
+    def __init__(self, **kwargs):
+        """Initializes the DistanceCalculation class for measuring object distances in video streams."""
+        super().__init__(**kwargs)
+
+        # Mouse event information
         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
+        self.centroids = []  # Initialize empty list to store centroids
 
     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.
+        Handles mouse events to select regions in a real-time video stream for distance calculation.
 
         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.
+            event (int): Type of mouse event (e.g., cv2.EVENT_MOUSEMOVE, cv2.EVENT_LBUTTONDOWN).
+            x (int): X-coordinate of the mouse pointer.
+            y (int): Y-coordinate of the mouse pointer.
+            flags (int): Flags associated with the event (e.g., cv2.EVENT_FLAG_CTRLKEY, cv2.EVENT_FLAG_SHIFTKEY).
+            param (Dict): Additional parameters passed to the function.
+
+        Examples:
+            >>> # Assuming 'dc' is an instance of DistanceCalculation
+            >>> cv2.setMouseCallback("window_name", dc.mouse_event_for_distance)
         """
-        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):
+                for box, track_id in zip(self.boxes, self.track_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:
+        elif 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):
+    def calculate(self, im0):
         """
-        Calculate the centroid of bounding box.
+        Processes a video frame and calculates the distance between two selected bounding boxes.
 
-        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.
+        This method extracts tracks from the input frame, annotates bounding boxes, and calculates the distance
+        between two user-selected objects if they have been chosen.
 
         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
+            im0 (numpy.ndarray): The input image frame to process.
 
-    def start_process(self, im0, tracks):
-        """
-        Calculate distance between two bounding boxes based on tracking data.
+        Returns:
+            (numpy.ndarray): The processed image frame with annotations and distance calculations.
 
-        Args:
-            im0 (nd array): Image
-            tracks (list): List of tracks obtained from the object tracking process.
+        Examples:
+            >>> import numpy as np
+            >>> from ultralytics.solutions import DistanceCalculation
+            >>> dc = DistanceCalculation()
+            >>> frame = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
+            >>> processed_frame = dc.calculate(frame)
         """
-        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(im0, line_width=self.line_width)  # Initialize annotator
+        self.extract_tracks(im0)  # Extract tracks
 
-        self.annotator = Annotator(self.im0, line_width=2)
-
-        for box, cls, track_id in zip(self.boxes, self.clss, self.trk_ids):
+        # Iterate over bounding boxes, track ids and classes index
+        for box, track_id, cls in zip(self.boxes, self.track_ids, self.clss):
             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():
+                for trk_id in self.selected_boxes.keys():
                     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
+            # Store user selected boxes in centroids list
+            self.centroids.extend(
+                [[int((box[0] + box[2]) // 2), int((box[1] + box[3]) // 2)] for box in self.selected_boxes.values()]
             )
+            # Calculate pixels distance
+            pixels_distance = math.sqrt(
+                (self.centroids[0][0] - self.centroids[1][0]) ** 2 + (self.centroids[0][1] - self.centroids[1][1]) ** 2
+            )
+            self.annotator.plot_distance_and_line(pixels_distance, self.centroids)
 
         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
-
+        self.display_output(im0)  # display output with base class function
+        cv2.setMouseCallback("Ultralytics Solutions", self.mouse_event_for_distance)
 
-if __name__ == "__main__":
-    DistanceCalculation()
+        return im0  # return output image for more usage

ultralytics/solutions/heatmap.py

@@ -1,281 +1,127 @@
-# Ultralytics YOLO 🚀, AGPL-3.0 license
-
-from collections import defaultdict
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 import cv2
 import numpy as np
 
-from ultralytics.utils.checks import check_imshow, check_requirements
+from ultralytics.solutions.object_counter import ObjectCounter
 from ultralytics.utils.plotting import Annotator
 
-check_requirements("shapely>=2.0.0")
-
-from shapely.geometry import LineString, Point, Polygon
-
 
-class Heatmap:
-    """A class to draw heatmaps in real-time video stream based on their tracks."""
+class Heatmap(ObjectCounter):
+    """
+    A class to draw heatmaps in real-time video streams based on object tracks.
 
-    def __init__(self):
-        """Initializes the heatmap class with default values for Visual, Image, track, count and heatmap parameters."""
+    This class extends the ObjectCounter class to generate and visualize heatmaps of object movements in video
+    streams. It uses tracked object positions to create a cumulative heatmap effect over time.
 
-        # Visual information
-        self.annotator = None
-        self.view_img = False
-        self.shape = "circle"
+    Attributes:
+        initialized (bool): Flag indicating whether the heatmap has been initialized.
+        colormap (int): OpenCV colormap used for heatmap visualization.
+        heatmap (np.ndarray): Array storing the cumulative heatmap data.
+        annotator (Annotator): Object for drawing annotations on the image.
 
-        # 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
+    Methods:
+        heatmap_effect: Calculates and updates the heatmap effect for a given bounding box.
+        generate_heatmap: Generates and applies the heatmap effect to each frame.
 
-        # Predict/track information
-        self.boxes = None
-        self.track_ids = None
-        self.clss = None
-        self.track_history = defaultdict(list)
+    Examples:
+        >>> from ultralytics.solutions import Heatmap
+        >>> heatmap = Heatmap(model="yolov8n.pt", colormap=cv2.COLORMAP_JET)
+        >>> frame = cv2.imread("frame.jpg")
+        >>> processed_frame = heatmap.generate_heatmap(frame)
+    """
 
-        # 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)
+    def __init__(self, **kwargs):
+        """Initializes the Heatmap class for real-time video stream heatmap generation based on object tracks."""
+        super().__init__(**kwargs)
 
-        # 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)
+        self.initialized = False  # bool variable for heatmap initialization
+        if self.region is not None:  # check if user provided the region coordinates
+            self.initialize_region()
 
-        # Decay factor
-        self.decay_factor = 0.99
-
-        # Check if environment support imshow
-        self.env_check = check_imshow(warn=True)
+        # store colormap
+        self.colormap = cv2.COLORMAP_PARULA if self.CFG["colormap"] is None else self.CFG["colormap"]
+        self.heatmap = None
 
-    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",
-    ):
+    def heatmap_effect(self, box):
         """
-        Configures the heatmap colormap, width, height and display parameters.
+        Efficiently calculates heatmap area and effect location for applying colormap.
 
         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)
+            box (List[float]): Bounding box coordinates [x0, y0, x1, y1].
 
-            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)
+        Examples:
+            >>> heatmap = Heatmap()
+            >>> box = [100, 100, 200, 200]
+            >>> heatmap.heatmap_effect(box)
+        """
+        x0, y0, x1, y1 = map(int, box)
+        radius_squared = (min(x1 - x0, y1 - y0) // 2) ** 2
 
-            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
+        # Create a meshgrid with region of interest (ROI) for vectorized distance calculations
+        xv, yv = np.meshgrid(np.arange(x0, x1), np.arange(y0, y1))
 
-        # Heatmap new frame
-        self.heatmap = np.zeros((int(self.imh), int(self.imw)), dtype=np.float32)
+        # Calculate squared distances from the center
+        dist_squared = (xv - ((x0 + x1) // 2)) ** 2 + (yv - ((y0 + y1) // 2)) ** 2
 
-        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
+        # Create a mask of points within the radius
+        within_radius = dist_squared <= radius_squared
 
-        # 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"
+        # Update only the values within the bounding box in a single vectorized operation
+        self.heatmap[y0:y1, x0:x1][within_radius] += 2
 
-    def extract_results(self, tracks):
+    def generate_heatmap(self, im0):
         """
-        Extracts results from the provided data.
+        Generate heatmap for each frame using Ultralytics.
 
         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()
+            im0 (np.ndarray): Input image array for processing.
 
-    def generate_heatmap(self, im0, tracks):
-        """
-        Generate heatmap based on tracking data.
+        Returns:
+            (np.ndarray): Processed image with heatmap overlay and object counts (if region is specified).
 
-        Args:
-            im0 (nd array): Image
-            tracks (list): List of tracks obtained from the object tracking process.
+        Examples:
+            >>> heatmap = Heatmap()
+            >>> im0 = cv2.imread("image.jpg")
+            >>> result = heatmap.generate_heatmap(im0)
         """
-        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
+        if not self.initialized:
+            self.heatmap = np.zeros_like(im0, dtype=np.float32) * 0.99
+        self.initialized = True  # Initialize heatmap only once
+
+        self.annotator = Annotator(im0, line_width=self.line_width)  # Initialize annotator
+        self.extract_tracks(im0)  # Extract tracks
+
+        # Iterate over bounding boxes, track ids and classes index
+        for box, track_id, cls in zip(self.boxes, self.track_ids, self.clss):
+            # Draw bounding box and counting region
+            self.heatmap_effect(box)
+
+            if self.region is not None:
+                self.annotator.draw_region(reg_pts=self.region, color=(104, 0, 123), thickness=self.line_width * 2)
+                self.store_tracking_history(track_id, box)  # Store track history
+                self.store_classwise_counts(cls)  # store classwise counts in dict
+                current_centroid = ((box[0] + box[2]) / 2, (box[1] + box[3]) / 2)
+                # Store tracking previous position and perform object counting
+                prev_position = None
+                if len(self.track_history[track_id]) > 1:
+                    prev_position = self.track_history[track_id][-2]
+                self.count_objects(current_centroid, track_id, prev_position, cls)  # Perform object counting
+
+        if self.region is not None:
+            self.display_counts(im0)  # Display the counts on the frame
 
         # 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,
+        if self.track_data.id is not None:
+            im0 = cv2.addWeighted(
+                im0,
+                0.5,
+                cv2.applyColorMap(
+                    cv2.normalize(self.heatmap, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8), self.colormap
+                ),
+                0.5,
+                0,
             )
 
-        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
-
-
-if __name__ == "__main__":
-    Heatmap()
+        self.display_output(im0)  # display output with base class function
+        return im0  # return output image for more usage