HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

ultralytics 8.3.16 PyTorch 2.5.0 support (#16998)

Browse source 
Signed-off-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: RizwanMunawar <chr043416@gmail.com>
Co-authored-by: Muhammad Rizwan Munawar <muhammadrizwanmunawar123@gmail.com>
This commit is contained in:
Glenn Jocher 2024-10-18 13:54:45 +02:00 • committed by GitHub
parent ef28f1078c
commit 8d7d1fe390
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
17 changed files with 570 additions and 144 deletions
Download patch file Download diff file Expand all files Collapse all files

64
ultralytics/solutions/heatmap.py
View file

@ -3,15 +3,40 @@
import cv2
import numpy as np
from ultralytics.solutions.object_counter import ObjectCounter # Import object counter class
from ultralytics.solutions.object_counter import ObjectCounter
from ultralytics.utils.plotting import Annotator
class Heatmap(ObjectCounter):
"""A class to draw heatmaps in real-time video stream based on their tracks."""
"""
A class to draw heatmaps in real-time video streams based on object tracks.
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.
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.
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.
Examples:
>>> from ultralytics.solutions import Heatmap
>>> heatmap = Heatmap(model="yolov8n.pt", colormap=cv2.COLORMAP_JET)
>>> results = heatmap("path/to/video.mp4")
>>> for result in results:
... print(result.speed) # Print inference speed
... cv2.imshow("Heatmap", result.plot())
... if cv2.waitKey(1) & 0xFF == ord("q"):
... break
"""
def __init__(self, **kwargs):
"""Initializes function for heatmap class with default values."""
"""Initializes the Heatmap class for real-time video stream heatmap generation based on object tracks."""
super().__init__(**kwargs)
self.initialized = False # bool variable for heatmap initialization
@ -23,10 +48,15 @@ class Heatmap(ObjectCounter):
def heatmap_effect(self, box):
"""
Efficient calculation of heatmap area and effect location for applying colormap.
Efficiently calculates heatmap area and effect location for applying colormap.
Args:
box (list): Bounding Box coordinates data [x0, y0, x1, y1]
box (List[float]): Bounding box coordinates [x0, y0, x1, y1].
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
@ -48,9 +78,15 @@ class Heatmap(ObjectCounter):
Generate heatmap for each frame using Ultralytics.
Args:
im0 (ndarray): Input image array for processing
im0 (np.ndarray): Input image array for processing.
Returns:
im0 (ndarray): Processed image for further usage
(np.ndarray): Processed image with heatmap overlay and object counts (if region is specified).
Examples:
>>> heatmap = Heatmap()
>>> im0 = cv2.imread("image.jpg")
>>> result = heatmap.generate_heatmap(im0)
"""
if not self.initialized:
self.heatmap = np.zeros_like(im0, dtype=np.float32) * 0.99
@ -70,16 +106,17 @@ class Heatmap(ObjectCounter):
self.store_classwise_counts(cls) # store classwise counts in dict
# Store tracking previous position and perform object counting
prev_position = self.track_history[track_id][-2] if len(self.track_history[track_id]) > 1 else None
prev_position = None
if len(self.track_history[track_id]) > 1:
prev_position = self.track_history[track_id][-2]
self.count_objects(self.track_line, box, track_id, prev_position, cls) # Perform object counting
self.display_counts(im0) if self.region is not None else None # Display the counts on the frame
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
im0 = (
im0
if self.track_data.id is None
else cv2.addWeighted(
if self.track_data.id is not None:
im0 = cv2.addWeighted(
im0,
0.5,
cv2.applyColorMap(
@ -88,7 +125,6 @@ class Heatmap(ObjectCounter):
0.5,
0,
)
)
self.display_output(im0) # display output with base class function
return im0 # return output image for more usage