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

ultralytics 8.0.239 Ultralytics Actions and hub-sdk adoption (#7431)

Browse source 
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Burhan <62214284+Burhan-Q@users.noreply.github.com>
Co-authored-by: Kayzwer <68285002+Kayzwer@users.noreply.github.com>
This commit is contained in:
Glenn Jocher 2024-01-10 03:16:08 +01:00 • committed by GitHub
parent e795277391
commit fe27db2f6e
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
139 changed files with 6870 additions and 5125 deletions
Download patch file Download diff file Expand all files Collapse all files

104
ultralytics/solutions/ai_gym.py
View file

@ -26,7 +26,7 @@ class AIGym:
self.angle = None
self.count = None
self.stage = None
self.pose_type = 'pushup'
self.pose_type = "pushup"
self.kpts_to_check = None
# Visual Information
@ -36,13 +36,15 @@ class AIGym:
# Check if environment support imshow
self.env_check = check_imshow(warn=True)
def set_args(self,
kpts_to_check,
line_thickness=2,
view_img=False,
pose_up_angle=145.0,
pose_down_angle=90.0,
pose_type='pullup'):
def set_args(
self,
kpts_to_check,
line_thickness=2,
view_img=False,
pose_up_angle=145.0,
pose_down_angle=90.0,
pose_type="pullup",
):
"""
Configures the AIGym line_thickness, save image and view image parameters
Args:
@ -72,65 +74,75 @@ class AIGym:
if frame_count == 1:
self.count = [0] * len(results[0])
self.angle = [0] * len(results[0])
self.stage = ['-' for _ in results[0]]
self.stage = ["-" for _ in results[0]]
self.keypoints = results[0].keypoints.data
self.annotator = Annotator(im0, line_width=2)
for ind, k in enumerate(reversed(self.keypoints)):
if self.pose_type == 'pushup' or self.pose_type == 'pullup':
self.angle[ind] = self.annotator.estimate_pose_angle(k[int(self.kpts_to_check[0])].cpu(),
k[int(self.kpts_to_check[1])].cpu(),
k[int(self.kpts_to_check[2])].cpu())
if self.pose_type == "pushup" or self.pose_type == "pullup":
self.angle[ind] = self.annotator.estimate_pose_angle(
k[int(self.kpts_to_check[0])].cpu(),
k[int(self.kpts_to_check[1])].cpu(),
k[int(self.kpts_to_check[2])].cpu(),
)
self.im0 = self.annotator.draw_specific_points(k, self.kpts_to_check, shape=(640, 640), radius=10)
if self.pose_type == 'abworkout':
self.angle[ind] = self.annotator.estimate_pose_angle(k[int(self.kpts_to_check[0])].cpu(),
k[int(self.kpts_to_check[1])].cpu(),
k[int(self.kpts_to_check[2])].cpu())
if self.pose_type == "abworkout":
self.angle[ind] = self.annotator.estimate_pose_angle(
k[int(self.kpts_to_check[0])].cpu(),
k[int(self.kpts_to_check[1])].cpu(),
k[int(self.kpts_to_check[2])].cpu(),
)
self.im0 = self.annotator.draw_specific_points(k, self.kpts_to_check, shape=(640, 640), radius=10)
if self.angle[ind] > self.poseup_angle:
self.stage[ind] = 'down'
if self.angle[ind] < self.posedown_angle and self.stage[ind] == 'down':
self.stage[ind] = 'up'
self.stage[ind] = "down"
if self.angle[ind] < self.posedown_angle and self.stage[ind] == "down":
self.stage[ind] = "up"
self.count[ind] += 1
self.annotator.plot_angle_and_count_and_stage(angle_text=self.angle[ind],
count_text=self.count[ind],
stage_text=self.stage[ind],
center_kpt=k[int(self.kpts_to_check[1])],
line_thickness=self.tf)
self.annotator.plot_angle_and_count_and_stage(
angle_text=self.angle[ind],
count_text=self.count[ind],
stage_text=self.stage[ind],
center_kpt=k[int(self.kpts_to_check[1])],
line_thickness=self.tf,
)
if self.pose_type == 'pushup':
if self.pose_type == "pushup":
if self.angle[ind] > self.poseup_angle:
self.stage[ind] = 'up'
if self.angle[ind] < self.posedown_angle and self.stage[ind] == 'up':
self.stage[ind] = 'down'
self.stage[ind] = "up"
if self.angle[ind] < self.posedown_angle and self.stage[ind] == "up":
self.stage[ind] = "down"
self.count[ind] += 1
self.annotator.plot_angle_and_count_and_stage(angle_text=self.angle[ind],
count_text=self.count[ind],
stage_text=self.stage[ind],
center_kpt=k[int(self.kpts_to_check[1])],
line_thickness=self.tf)
if self.pose_type == 'pullup':
self.annotator.plot_angle_and_count_and_stage(
angle_text=self.angle[ind],
count_text=self.count[ind],
stage_text=self.stage[ind],
center_kpt=k[int(self.kpts_to_check[1])],
line_thickness=self.tf,
)
if self.pose_type == "pullup":
if self.angle[ind] > self.poseup_angle:
self.stage[ind] = 'down'
if self.angle[ind] < self.posedown_angle and self.stage[ind] == 'down':
self.stage[ind] = 'up'
self.stage[ind] = "down"
if self.angle[ind] < self.posedown_angle and self.stage[ind] == "down":
self.stage[ind] = "up"
self.count[ind] += 1
self.annotator.plot_angle_and_count_and_stage(angle_text=self.angle[ind],
count_text=self.count[ind],
stage_text=self.stage[ind],
center_kpt=k[int(self.kpts_to_check[1])],
line_thickness=self.tf)
self.annotator.plot_angle_and_count_and_stage(
angle_text=self.angle[ind],
count_text=self.count[ind],
stage_text=self.stage[ind],
center_kpt=k[int(self.kpts_to_check[1])],
line_thickness=self.tf,
)
self.annotator.kpts(k, shape=(640, 640), radius=1, kpt_line=True)
if self.env_check and self.view_img:
cv2.imshow('Ultralytics YOLOv8 AI GYM', self.im0)
if cv2.waitKey(1) & 0xFF == ord('q'):
cv2.imshow("Ultralytics YOLOv8 AI GYM", self.im0)
if cv2.waitKey(1) & 0xFF == ord("q"):
return
return self.im0
if __name__ == '__main__':
if __name__ == "__main__":
AIGym()

31
ultralytics/solutions/distance_calculation.py
View file

@ -41,13 +41,15 @@ class DistanceCalculation:
# 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)):
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.
@ -129,8 +131,9 @@ class DistanceCalculation:
distance (float): Distance between two centroids
"""
cv2.rectangle(self.im0, (15, 25), (280, 70), (255, 255, 255), -1)
cv2.putText(self.im0, f'Distance : {distance:.2f}m', (20, 55), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 0), 2,
cv2.LINE_AA)
cv2.putText(
self.im0, f"Distance : {distance:.2f}m", (20, 55), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 0), 2, cv2.LINE_AA
)
cv2.line(self.im0, self.centroids[0], self.centroids[1], self.line_color, 3)
cv2.circle(self.im0, self.centroids[0], 6, self.centroid_color, -1)
cv2.circle(self.im0, self.centroids[1], 6, self.centroid_color, -1)
@ -179,13 +182,13 @@ class DistanceCalculation:
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)
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'):
if cv2.waitKey(1) & 0xFF == ord("q"):
return
if __name__ == '__main__':
if __name__ == "__main__":
DistanceCalculation()

112
ultralytics/solutions/heatmap.py
View file

@ -8,7 +8,7 @@ import numpy as np
from ultralytics.utils.checks import check_imshow, check_requirements
from ultralytics.utils.plotting import Annotator
check_requirements('shapely>=2.0.0')
check_requirements("shapely>=2.0.0")
from shapely.geometry import LineString, Point, Polygon
@ -22,7 +22,7 @@ class Heatmap:
# Visual information
self.annotator = None
self.view_img = False
self.shape = 'circle'
self.shape = "circle"
# Image information
self.imw = None
@ -63,23 +63,25 @@ class Heatmap:
# Check if environment support imshow
self.env_check = check_imshow(warn=True)
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 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",
):
"""
Configures the heatmap colormap, width, height and display parameters.
@ -111,20 +113,19 @@ class Heatmap:
# Region and line selection
if count_reg_pts is not None:
if len(count_reg_pts) == 2:
print('Line Counter Initiated.')
print("Line Counter Initiated.")
self.count_reg_pts = count_reg_pts
self.counting_region = LineString(count_reg_pts)
elif len(count_reg_pts) == 4:
print('Region Counter Initiated.')
print("Region Counter Initiated.")
self.count_reg_pts = count_reg_pts
self.counting_region = Polygon(self.count_reg_pts)
else:
print('Region or line points Invalid, 2 or 4 points supported')
print('Using Line Counter Now')
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
# Heatmap new frame
@ -140,10 +141,10 @@ class Heatmap:
self.shape = shape
# shape of heatmap, if not selected
if self.shape not in ['circle', 'rect']:
if self.shape not in ["circle", "rect"]:
print("Unknown shape value provided, 'circle' & 'rect' supported")
print('Using Circular shape now')
self.shape = 'circle'
print("Using Circular shape now")
self.shape = "circle"
def extract_results(self, tracks):
"""
@ -177,27 +178,26 @@ class Heatmap:
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)
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':
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
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])])
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
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]
@ -226,26 +226,26 @@ class Heatmap:
self.in_counts += 1
else:
for box, cls in zip(self.boxes, self.clss):
if self.shape == 'circle':
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
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])])
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
self.heatmap[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])] += 2
# 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 = 'In Count : ' + f'{self.in_counts}'
outcount_label = 'OutCount : ' + f'{self.out_counts}'
incount_label = "In Count : " + f"{self.in_counts}"
outcount_label = "OutCount : " + f"{self.out_counts}"
# Display counts based on user choice
counts_label = None
@ -256,13 +256,15 @@ class Heatmap:
elif not self.view_out_counts:
counts_label = incount_label
else:
counts_label = incount_label + ' ' + outcount_label
counts_label = 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)
self.annotator.count_labels(
counts=counts_label,
count_txt_size=self.count_txt_thickness,
txt_color=self.count_txt_color,
color=self.count_color,
)
self.im0 = cv2.addWeighted(self.im0, 1 - self.heatmap_alpha, heatmap_colored, self.heatmap_alpha, 0)
@ -273,11 +275,11 @@ class Heatmap:
def display_frames(self):
"""Display frame."""
cv2.imshow('Ultralytics Heatmap', self.im0)
cv2.imshow("Ultralytics Heatmap", self.im0)
if cv2.waitKey(1) & 0xFF == ord('q'):
if cv2.waitKey(1) & 0xFF == ord("q"):
return
if __name__ == '__main__':
if __name__ == "__main__":
Heatmap()

84
ultralytics/solutions/object_counter.py
View file

@ -7,7 +7,7 @@ import cv2
from ultralytics.utils.checks import check_imshow, check_requirements
from ultralytics.utils.plotting import Annotator, colors
check_requirements('shapely>=2.0.0')
check_requirements("shapely>=2.0.0")
from shapely.geometry import LineString, Point, Polygon
@ -56,22 +56,24 @@ class ObjectCounter:
# Check if environment support imshow
self.env_check = check_imshow(warn=True)
def set_args(self,
classes_names,
reg_pts,
count_reg_color=(255, 0, 255),
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),
count_color=(255, 255, 255),
track_color=(0, 255, 0),
region_thickness=5,
line_dist_thresh=15):
def set_args(
self,
classes_names,
reg_pts,
count_reg_color=(255, 0, 255),
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),
count_color=(255, 255, 255),
track_color=(0, 255, 0),
region_thickness=5,
line_dist_thresh=15,
):
"""
Configures the Counter's image, bounding box line thickness, and counting region points.
@ -101,16 +103,16 @@ class ObjectCounter:
# Region and line selection
if len(reg_pts) == 2:
print('Line Counter Initiated.')
print("Line Counter Initiated.")
self.reg_pts = reg_pts
self.counting_region = LineString(self.reg_pts)
elif len(reg_pts) == 4:
print('Region Counter Initiated.')
print("Region Counter Initiated.")
self.reg_pts = reg_pts
self.counting_region = Polygon(self.reg_pts)
else:
print('Invalid Region points provided, region_points can be 2 or 4')
print('Using Line Counter Now')
print("Invalid Region points provided, region_points can be 2 or 4")
print("Using Line Counter Now")
self.counting_region = LineString(self.reg_pts)
self.names = classes_names
@ -164,8 +166,9 @@ class ObjectCounter:
# Extract tracks
for box, track_id, cls in zip(boxes, track_ids, clss):
self.annotator.box_label(box, label=str(track_id) + ':' + self.names[cls],
color=colors(int(cls), True)) # Draw bounding box
self.annotator.box_label(
box, label=str(track_id) + ":" + self.names[cls], color=colors(int(cls), True)
) # Draw bounding box
# Draw Tracks
track_line = self.track_history[track_id]
@ -175,9 +178,9 @@ class ObjectCounter:
# Draw track trails
if self.draw_tracks:
self.annotator.draw_centroid_and_tracks(track_line,
color=self.track_color,
track_thickness=self.track_thickness)
self.annotator.draw_centroid_and_tracks(
track_line, color=self.track_color, track_thickness=self.track_thickness
)
# Count objects
if len(self.reg_pts) == 4:
@ -199,8 +202,8 @@ class ObjectCounter:
else:
self.in_counts += 1
incount_label = 'In Count : ' + f'{self.in_counts}'
outcount_label = 'OutCount : ' + f'{self.out_counts}'
incount_label = "In Count : " + f"{self.in_counts}"
outcount_label = "OutCount : " + f"{self.out_counts}"
# Display counts based on user choice
counts_label = None
@ -211,24 +214,27 @@ class ObjectCounter:
elif not self.view_out_counts:
counts_label = incount_label
else:
counts_label = incount_label + ' ' + outcount_label
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)
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."""
if self.env_check:
cv2.namedWindow('Ultralytics YOLOv8 Object Counter')
cv2.namedWindow("Ultralytics YOLOv8 Object Counter")
if len(self.reg_pts) == 4: # only add mouse event If user drawn region
cv2.setMouseCallback('Ultralytics YOLOv8 Object Counter', self.mouse_event_for_region,
{'region_points': self.reg_pts})
cv2.imshow('Ultralytics YOLOv8 Object Counter', self.im0)
cv2.setMouseCallback(
"Ultralytics YOLOv8 Object Counter", self.mouse_event_for_region, {"region_points": self.reg_pts}
)
cv2.imshow("Ultralytics YOLOv8 Object Counter", self.im0)
# Break Window
if cv2.waitKey(1) & 0xFF == ord('q'):
if cv2.waitKey(1) & 0xFF == ord("q"):
return
def start_counting(self, im0, tracks):
@ -254,5 +260,5 @@ class ObjectCounter:
return self.im0
if __name__ == '__main__':
if __name__ == "__main__":
ObjectCounter()

29
ultralytics/solutions/speed_estimation.py
View file

@ -66,7 +66,7 @@ class SpeedEstimator:
spdl_dist_thresh (int): Euclidean distance threshold for speed line
"""
if reg_pts is None:
print('Region points not provided, using default values')
print("Region points not provided, using default values")
else:
self.reg_pts = reg_pts
self.names = names
@ -114,8 +114,9 @@ class SpeedEstimator:
cls (str): object class name
track (list): tracking history for tracks path drawing
"""
speed_label = str(int(
self.dist_data[track_id])) + 'km/ph' if track_id in self.dist_data else self.names[int(cls)]
speed_label = (
str(int(self.dist_data[track_id])) + "km/ph" if track_id in self.dist_data else self.names[int(cls)]
)
bbox_color = colors(int(track_id)) if track_id in self.dist_data else (255, 0, 255)
self.annotator.box_label(box, speed_label, bbox_color)
@ -132,19 +133,16 @@ class SpeedEstimator:
"""
if self.reg_pts[0][0] < track[-1][0] < self.reg_pts[1][0]:
if self.reg_pts[1][1] - self.spdl_dist_thresh < track[-1][1] < self.reg_pts[1][1] + self.spdl_dist_thresh:
direction = "known"
if (self.reg_pts[1][1] - self.spdl_dist_thresh < track[-1][1] < self.reg_pts[1][1] + self.spdl_dist_thresh):
direction = 'known'
elif (self.reg_pts[0][1] - self.spdl_dist_thresh < track[-1][1] <
self.reg_pts[0][1] + self.spdl_dist_thresh):
direction = 'known'
elif self.reg_pts[0][1] - self.spdl_dist_thresh < track[-1][1] < self.reg_pts[0][1] + self.spdl_dist_thresh:
direction = "known"
else:
direction = 'unknown'
if self.trk_previous_times[trk_id] != 0 and direction != 'unknown':
direction = "unknown"
if self.trk_previous_times[trk_id] != 0 and direction != "unknown":
if trk_id not in self.trk_idslist:
self.trk_idslist.append(trk_id)
@ -178,7 +176,6 @@ class SpeedEstimator:
self.annotator.draw_region(reg_pts=self.reg_pts, color=(255, 0, 0), thickness=self.region_thickness)
for box, trk_id, cls in zip(self.boxes, self.trk_ids, self.clss):
track = self.store_track_info(trk_id, box)
if trk_id not in self.trk_previous_times:
@ -194,10 +191,10 @@ class SpeedEstimator:
def display_frames(self):
"""Display frame."""
cv2.imshow('Ultralytics Speed Estimation', self.im0)
if cv2.waitKey(1) & 0xFF == ord('q'):
cv2.imshow("Ultralytics Speed Estimation", self.im0)
if cv2.waitKey(1) & 0xFF == ord("q"):
return
if __name__ == '__main__':
if __name__ == "__main__":
SpeedEstimator()