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.

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

Solutions updates #9548

Merged
merged 2 commits into from
Apr 5, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 20 additions & 21 deletions docs/en/guides/heatmaps.md
Original file line number Diff line number Diff line change
Expand Up @@ -293,27 +293,26 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult

### Arguments `set_args`

| Name | Type | Default | Description |
|-----------------------|----------------|---------------------|-----------------------------------------------------------|
| `view_img` | `bool` | `False` | Display the frame with heatmap |
| `colormap` | `cv2.COLORMAP` | `None` | cv2.COLORMAP for heatmap |
| `imw` | `int` | `None` | Width of Heatmap |
| `imh` | `int` | `None` | Height of Heatmap |
| `view_in_counts` | `bool` | `True` | Display in-counts only on video frame |
| `view_out_counts` | `bool` | `True` | Display out-counts only on video frame |
| `classes_names` | `dict` | `model.model.names` | Dictionary of Class Names |
| `heatmap_alpha` | `float` | `0.5` | Heatmap alpha value |
| `count_reg_pts` | `list` | `None` | Object counting region points |
| `count_txt_thickness` | `int` | `2` | Count values text size |
| `count_txt_color` | `RGB Color` | `(0, 0, 0)` | Foreground color for Object counts text |
| `count_reg_color` | `RGB Color` | `(255, 0, 255)` | Counting region color |
| `region_thickness` | `int` | `5` | Counting region thickness value |
| `decay_factor` | `float` | `0.99` | Decay factor for heatmap area removal after specific time |
| `shape` | `str` | `circle` | Heatmap shape for display "rect" or "circle" supported |
| `line_dist_thresh` | `int` | `15` | Euclidean Distance threshold for line counter |
| `fontsize` | `float` | `0.6` | Font size of counting text |
| `line_color` | `RGB Color` | `(255, 255, 255)` | Count highlighter color |
| `cls_txtdisplay_gap` | `int` | `50` | Display gap between each class count |
| Name | Type | Default | Description |
|----------------------|----------------|---------------------|-----------------------------------------------------------|
| `view_img` | `bool` | `False` | Display the frame with heatmap |
| `colormap` | `cv2.COLORMAP` | `None` | cv2.COLORMAP for heatmap |
| `imw` | `int` | `None` | Width of Heatmap |
| `imh` | `int` | `None` | Height of Heatmap |
| `line_thickness` | `int` | `2` | Increase bounding boxes and count text thickness |
| `view_in_counts` | `bool` | `True` | Display in-counts only on video frame |
| `view_out_counts` | `bool` | `True` | Display out-counts only on video frame |
| `classes_names` | `dict` | `model.model.names` | Dictionary of Class Names |
| `heatmap_alpha` | `float` | `0.5` | Heatmap alpha value |
| `count_reg_pts` | `list` | `None` | Object counting region points |
| `count_txt_color` | `RGB Color` | `(0, 0, 0)` | Foreground color for Object counts text |
| `count_reg_color` | `RGB Color` | `(255, 0, 255)` | Counting region color |
| `region_thickness` | `int` | `5` | Counting region thickness value |
| `decay_factor` | `float` | `0.99` | Decay factor for heatmap area removal after specific time |
| `shape` | `str` | `circle` | Heatmap shape for display "rect" or "circle" supported |
| `line_dist_thresh` | `int` | `15` | Euclidean Distance threshold for line counter |
| `count_bg_color` | `RGB Color` | `(255, 255, 255)` | Count highlighter color |
| `cls_txtdisplay_gap` | `int` | `50` | Display gap between each class count |

### Arguments `model.track`

Expand Down
47 changes: 24 additions & 23 deletions docs/en/guides/object-counting.md
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,8 @@ Object counting with [Ultralytics YOLOv8](https://github.com/ultralytics/ultraly
counter.set_args(view_img=True,
reg_pts=region_points,
classes_names=model.names,
draw_tracks=True)
draw_tracks=True,
line_thickness=2)

while cap.isOpened():
success, im0 = cap.read()
Expand Down Expand Up @@ -105,7 +106,8 @@ Object counting with [Ultralytics YOLOv8](https://github.com/ultralytics/ultraly
counter.set_args(view_img=True,
reg_pts=region_points,
classes_names=model.names,
draw_tracks=True)
draw_tracks=True,
line_thickness=2)

while cap.isOpened():
success, im0 = cap.read()
Expand Down Expand Up @@ -148,7 +150,8 @@ Object counting with [Ultralytics YOLOv8](https://github.com/ultralytics/ultraly
counter.set_args(view_img=True,
reg_pts=line_points,
classes_names=model.names,
draw_tracks=True)
draw_tracks=True,
line_thickness=2)

while cap.isOpened():
success, im0 = cap.read()
Expand Down Expand Up @@ -191,7 +194,8 @@ Object counting with [Ultralytics YOLOv8](https://github.com/ultralytics/ultraly
counter.set_args(view_img=True,
reg_pts=line_points,
classes_names=model.names,
draw_tracks=True)
draw_tracks=True,
line_thickness=2)

while cap.isOpened():
success, im0 = cap.read()
Expand All @@ -215,25 +219,22 @@ Object counting with [Ultralytics YOLOv8](https://github.com/ultralytics/ultraly

### Optional Arguments `set_args`

| Name | Type | Default | Description |
|-----------------------|-------------|----------------------------|-----------------------------------------------|
| `view_img` | `bool` | `False` | Display frames with counts |
| `view_in_counts` | `bool` | `True` | Display in-counts only on video frame |
| `view_out_counts` | `bool` | `True` | Display out-counts only on video frame |
| `line_thickness` | `int` | `2` | Increase bounding boxes thickness |
| `reg_pts` | `list` | `[(20, 400), (1260, 400)]` | Points defining the Region Area |
| `classes_names` | `dict` | `model.model.names` | Dictionary of Class Names |
| `region_color` | `RGB Color` | `(255, 0, 255)` | Color of the Object counting Region or Line |
| `track_thickness` | `int` | `2` | Thickness of Tracking Lines |
| `draw_tracks` | `bool` | `False` | Enable drawing Track lines |
| `track_color` | `RGB Color` | `(0, 255, 0)` | Color for each track line |
| `line_dist_thresh` | `int` | `15` | Euclidean Distance threshold for line counter |
| `count_txt_thickness` | `int` | `2` | Thickness of Object counts text |
| `count_txt_color` | `RGB Color` | `(255, 255, 255)` | Foreground color for Object counts text |
| `region_thickness` | `int` | `5` | Thickness for object counter region or line |
| `fontsize` | `float` | `0.6` | Font size of counting text |
| `line_color` | `RGB Color` | `(255, 255, 255)` | Count highlighter color |
| `cls_txtdisplay_gap` | `int` | `50` | Display gap between each class count |
| Name | Type | Default | Description |
|-----------------------|-------------|----------------------------|--------------------------------------------------|
| `view_img` | `bool` | `False` | Display frames with counts |
| `view_in_counts` | `bool` | `True` | Display in-counts only on video frame |
| `view_out_counts` | `bool` | `True` | Display out-counts only on video frame |
| `line_thickness` | `int` | `2` | Increase bounding boxes and count text thickness |
| `reg_pts` | `list` | `[(20, 400), (1260, 400)]` | Points defining the Region Area |
| `classes_names` | `dict` | `model.model.names` | Dictionary of Class Names |
| `count_reg_color` | `RGB Color` | `(255, 0, 255)` | Color of the Object counting Region or Line |
| `track_thickness` | `int` | `2` | Thickness of Tracking Lines |
| `draw_tracks` | `bool` | `False` | Enable drawing Track lines |
| `track_color` | `RGB Color` | `(0, 255, 0)` | Color for each track line |
| `line_dist_thresh` | `int` | `15` | Euclidean Distance threshold for line counter |
| `count_txt_color` | `RGB Color` | `(255, 255, 255)` | Foreground color for Object counts text |
| `region_thickness` | `int` | `5` | Thickness for object counter region or line |
| `count_bg_color` | `RGB Color` | `(255, 255, 255)` | Count highlighter color |

### Arguments `model.track`

Expand Down
41 changes: 14 additions & 27 deletions ultralytics/solutions/heatmap.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ def __init__(self):
self.imw = None
self.imh = None
self.im0 = None
self.tf = 2
self.view_in_counts = True
self.view_out_counts = True

Expand All @@ -56,11 +57,9 @@ def __init__(self):
self.out_counts = 0
self.count_ids = []
self.class_wise_count = {}
self.count_txt_thickness = 0
self.count_txt_color = (255, 255, 255)
self.line_color = (255, 255, 255)
self.count_txt_color = (0, 0, 0)
self.count_bg_color = (255, 255, 255)
self.cls_txtdisplay_gap = 50
self.fontsize = 0.6

# Decay factor
self.decay_factor = 0.99
Expand All @@ -79,16 +78,14 @@ def set_args(
view_in_counts=True,
view_out_counts=True,
count_reg_pts=None,
count_txt_thickness=2,
count_txt_color=(255, 255, 255),
fontsize=0.8,
line_color=(255, 255, 255),
count_txt_color=(0, 0, 0),
count_bg_color=(255, 255, 255),
count_reg_color=(255, 0, 255),
region_thickness=5,
line_dist_thresh=15,
line_thickness=2,
decay_factor=0.99,
shape="circle",
cls_txtdisplay_gap=50,
):
"""
Configures the heatmap colormap, width, height and display parameters.
Expand All @@ -98,22 +95,21 @@ def set_args(
imw (int): The width of the frame.
imh (int): The height of the frame.
classes_names (dict): Classes names
line_thickness (int): Line thickness for bounding boxes.
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
fontsize (float): Text display font size
line_color (RGB color): count highlighter line color
count_bg_color (RGB color): count highlighter line color
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
cls_txtdisplay_gap (int): Display gap between each class count
"""
self.tf = line_thickness
self.names = classes_names
self.imw = imw
self.imh = imh
Expand Down Expand Up @@ -141,16 +137,13 @@ def set_args(
# 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.fontsize = fontsize
self.line_color = line_color
self.count_bg_color = count_bg_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
self.cls_txtdisplay_gap = cls_txtdisplay_gap

# shape of heatmap, if not selected
if self.shape not in {"circle", "rect"}:
Expand Down Expand Up @@ -185,7 +178,7 @@ def generate_heatmap(self, im0, tracks):
return im0
self.heatmap *= self.decay_factor # decay factor
self.extract_results(tracks)
self.annotator = Annotator(self.im0, self.count_txt_thickness, None)
self.annotator = Annotator(self.im0, self.tf, None)

if self.count_reg_pts is not None:
# Draw counting region
Expand Down Expand Up @@ -239,11 +232,8 @@ def generate_heatmap(self, im0, tracks):

# Count objects using line
elif len(self.count_reg_pts) == 2:
is_inside = (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0

if prev_position is not None and is_inside and track_id not in self.count_ids:
if prev_position is not None and track_id not in self.count_ids:
distance = Point(track_line[-1]).distance(self.counting_region)

if distance < self.line_dist_thresh and track_id not in self.count_ids:
self.count_ids.append(track_id)

Expand Down Expand Up @@ -293,11 +283,8 @@ def generate_heatmap(self, im0, tracks):
if self.count_reg_pts is not None and label is not None:
self.annotator.display_counts(
counts=label,
tf=self.count_txt_thickness,
fontScale=self.fontsize,
txt_color=self.count_txt_color,
line_color=self.line_color,
classwise_txtgap=self.cls_txtdisplay_gap,
count_txt_color=self.count_txt_color,
count_bg_color=self.count_bg_color,
)

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