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

ultralytics 8.0.237 cv2.CAP_PROP fix and in_counts and out_counts displays (#7380)

Browse source 
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: ayush chaurasia <ayush.chaurarsia@gmail.com>
Co-authored-by: Muhammad Rizwan Munawar <chr043416@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: 曾逸夫（Zeng Yifu） <41098760+Zengyf-CVer@users.noreply.github.com>
This commit is contained in:
Glenn Jocher 2024-01-07 19:06:57 +01:00 • committed by GitHub
parent 71fe5e919d
commit 8c2b2f56b2
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
15 changed files with 108 additions and 90 deletions
Download patch file Download diff file Expand all files Collapse all files

2
ultralytics/__init__.py
View file

@ -1,6 +1,6 @@
# Ultralytics YOLO 🚀, AGPL-3.0 license
__version__ = '8.0.236'
__version__ = '8.0.237'
from ultralytics.data.explorer.explorer import Explorer
from ultralytics.models import RTDETR, SAM, YOLO

2
ultralytics/solutions/heatmap.py
View file

@ -120,7 +120,7 @@ class Heatmap:
self.counting_region = Polygon([(20, 400), (1260, 400)]) # dummy points
# Heatmap new frame
self.heatmap = np.zeros((int(self.imw), int(self.imh)), dtype=np.float32)
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

30
ultralytics/solutions/object_counter.py
View file

@ -33,6 +33,8 @@ class ObjectCounter:
self.im0 = None
self.tf = None
self.view_img = False
self.view_in_counts = True
self.view_out_counts = True
self.names = None # Classes names
self.annotator = None # Annotator
@ -61,6 +63,8 @@ class ObjectCounter:
line_thickness=2,
track_thickness=2,
view_img=False,
view_in_counts=True,
view_out_counts=True,
draw_tracks=False,
count_txt_thickness=2,
count_txt_color=(0, 0, 0),
@ -74,6 +78,8 @@ class ObjectCounter:
Args:
line_thickness (int): Line thickness for bounding boxes.
view_img (bool): Flag to control whether to display the video stream.
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.
reg_pts (list): Initial list of points defining the counting region.
classes_names (dict): Classes names
track_thickness (int): Track thickness
@ -88,6 +94,8 @@ class ObjectCounter:
"""
self.tf = line_thickness
self.view_img = view_img
self.view_in_counts = view_in_counts
self.view_out_counts = view_out_counts
self.track_thickness = track_thickness
self.draw_tracks = draw_tracks
@ -192,11 +200,23 @@ class ObjectCounter:
incount_label = 'In Count : ' + f'{self.in_counts}'
outcount_label = 'OutCount : ' + f'{self.out_counts}'
self.annotator.count_labels(in_count=incount_label,
out_count=outcount_label,
count_txt_size=self.count_txt_thickness,
txt_color=self.count_txt_color,
color=self.count_color)
# 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 = incount_label + ' ' + outcount_label
if 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)
def display_frames(self):
"""Display frame."""

32
ultralytics/utils/plotting.py
View file

@ -291,12 +291,11 @@ class Annotator:
cv2.polylines(self.im, [points], isClosed=False, color=color, thickness=track_thickness)
cv2.circle(self.im, (int(track[-1][0]), int(track[-1][1])), track_thickness * 2, color, -1)
def count_labels(self, in_count=0, out_count=0, count_txt_size=2, color=(255, 255, 255), txt_color=(0, 0, 0)):
def count_labels(self, counts=0, count_txt_size=2, color=(255, 255, 255), txt_color=(0, 0, 0)):
"""
Plot counts for object counter
Args:
in_count (int): in count value
out_count (int): out count value
counts (int): objects counts value
count_txt_size (int): text size for counts display
color (tuple): background color of counts display
txt_color (tuple): text color of counts display
@ -307,37 +306,24 @@ class Annotator:
gap = int(24 * tl) # gap between in_count and out_count based on line_thickness
# Get text size for in_count and out_count
t_size_in = cv2.getTextSize(str(in_count), 0, fontScale=tl / 2, thickness=tf)[0]
t_size_out = cv2.getTextSize(str(out_count), 0, fontScale=tl / 2, thickness=tf)[0]
t_size_in = cv2.getTextSize(str(counts), 0, fontScale=tl / 2, thickness=tf)[0]
# Calculate positions for in_count and out_count labels
text_width = max(t_size_in[0], t_size_out[0])
text_x1 = (self.im.shape[1] - text_width - 120 * self.tf) // 2 - gap
text_x2 = (self.im.shape[1] - text_width + 120 * self.tf) // 2 + gap
text_y = max(t_size_in[1], t_size_out[1])
# Calculate positions for counts label
text_width = t_size_in[0]
text_x = (self.im.shape[1] - text_width) // 2 # Center x-coordinate
text_y = t_size_in[1]
# Create a rounded rectangle for in_count
cv2.rectangle(self.im, (text_x1 - 5, text_y - 5), (text_x1 + text_width + 7, text_y + t_size_in[1] + 7), color,
cv2.rectangle(self.im, (text_x - 5, text_y - 5), (text_x + text_width + 7, text_y + t_size_in[1] + 7), color,
-1)
cv2.putText(self.im,
str(in_count), (text_x1, text_y + t_size_in[1]),
str(counts), (text_x, text_y + t_size_in[1]),
0,
tl / 2,
txt_color,
self.tf,
lineType=cv2.LINE_AA)
# Create a rounded rectangle for out_count
cv2.rectangle(self.im, (text_x2 - 5, text_y - 5), (text_x2 + text_width + 7, text_y + t_size_out[1] + 7), color,
-1)
cv2.putText(self.im,
str(out_count), (text_x2, text_y + t_size_out[1]),
0,
tl / 2,
txt_color,
thickness=self.tf,
lineType=cv2.LINE_AA)
@staticmethod
def estimate_pose_angle(a, b, c):
"""Calculate the pose angle for object