Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Added homography mappings to heatmap annotator #1057

Open
wants to merge 1 commit into
base: develop
Choose a base branch
from
+35 −2
Open

Added homography mappings to heatmap annotator #1057

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
37 changes: 35 additions & 2 deletions supervision/annotators/core.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1274,7 +1274,9 @@ class HeatMapAnnotator:
"""
A class for drawing heatmaps on an image based on provided detections.
Heat accumulates over time and is drawn as a semi-transparent overlay
of blurred circles.
of blurred circles. If homography mappings are provided the annotations
will be drawn on the transformed coordinates, typically used to transform
the image angle into a top-down view used in floor plans
"""

def __init__(
Expand All @@ -1285,6 +1287,7 @@ def __init__(
kernel_size: int = 25,
top_hue: int = 0,
low_hue: int = 125,
mappings: list = None,
):
"""
Args:
Expand All @@ -1295,6 +1298,11 @@ def __init__(
kernel_size (int): Kernel size for blurring the heatmap.
top_hue (int): Hue at the top of the heatmap. Defaults to 0 (red).
low_hue (int): Hue at the bottom of the heatmap. Defaults to 125 (blue).
mappings (list): Contains two lists, one for source and one for destination
xy coordinates. Used to perform homography mappings.
Must be equal number of coordinates in each list and at
least four points are required for successful
transformation
"""
self.position = position
self.opacity = opacity
Expand All @@ -1303,6 +1311,18 @@ def __init__(
self.heat_mask = None
self.top_hue = top_hue
self.low_hue = low_hue
self.mappings = mappings

# calculate homography matrix
if self.mappings:
pts_src = np.array(self.mappings[0])
pts_dst = np.array(self.mappings[1])
# Check if both pts_src and pts_dst have at least 4 points
if len(pts_src) < 4 or len(pts_dst) < 4:
raise ValueError(
"Require four or more mapping points for homography mappings."
)
self.homography_matrix, status = cv2.findHomography(pts_src, pts_dst)

@scene_to_annotator_img_type
def annotate(self, scene: ImageType, detections: Detections) -> ImageType:
Expand All @@ -1319,6 +1339,10 @@ def annotate(self, scene: ImageType, detections: Detections) -> ImageType:
The annotated image, matching the type of `scene` (`numpy.ndarray`
or `PIL.Image.Image`)

Raises:
Value error if less than four pairs of xy coordinates are provided in
the mappings list for both the source and destination scenes

Example:
```python
import supervision as sv
Expand Down Expand Up @@ -1349,7 +1373,16 @@ def annotate(self, scene: ImageType, detections: Detections) -> ImageType:
self.heat_mask = np.zeros(scene.shape[:2])
mask = np.zeros(scene.shape[:2])
for xy in detections.get_anchors_coordinates(self.position):
cv2.circle(mask, (int(xy[0]), int(xy[1])), self.radius, 1, -1)
# calculate transform based on homography mappings if provided
if self.mappings:
a = np.array([xy], dtype="float32")
a = np.array([a])
xy = cv2.perspectiveTransform(a, self.homography_matrix)
cv2.circle(
mask, (int(xy[0][0][0]), int(xy[0][0][1])), self.radius, 1, -1
)
else:
cv2.circle(mask, (int(xy[0]), int(xy[1])), self.radius, 1, -1)
self.heat_mask = mask + self.heat_mask
temp = self.heat_mask.copy()
temp = self.low_hue - temp / temp.max() * (self.low_hue - self.top_hue)
Expand Down