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

Python refactorings and simplifications (#7549)

Browse source 
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: Hassaan Farooq <103611273+hassaanfarooq01@users.noreply.github.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
This commit is contained in:
Glenn Jocher 2024-01-12 19:34:03 +01:00 • committed by GitHub
parent 0da13831cf
commit f6309b8e70
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
40 changed files with 127 additions and 189 deletions
Download patch file Download diff file Expand all files Collapse all files

2
ultralytics/solutions/ai_gym.py
View file

@ -79,7 +79,7 @@ class AIGym:
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":
if self.pose_type in ["pushup", "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(),

12
ultralytics/solutions/distance_calculation.py
View file

@ -86,10 +86,9 @@ class DistanceCalculation:
self.left_mouse_count += 1
if self.left_mouse_count <= 2:
for box, track_id in zip(self.boxes, self.trk_ids):
if box[0] < x < box[2] and box[1] < y < box[3]:
if track_id not in self.selected_boxes:
self.selected_boxes[track_id] = []
self.selected_boxes[track_id] = box
if box[0] < x < box[2] and box[1] < y < box[3] and track_id not in self.selected_boxes:
self.selected_boxes[track_id] = []
self.selected_boxes[track_id] = box
if event == cv2.EVENT_RBUTTONDOWN:
self.selected_boxes = {}
@ -149,10 +148,7 @@ class DistanceCalculation:
if tracks[0].boxes.id is None:
if self.view_img:
self.display_frames()
return
else:
return
return
self.extract_tracks(tracks)
self.annotator = Annotator(self.im0, line_width=2)

37
ultralytics/solutions/heatmap.py
View file

@ -169,10 +169,7 @@ class Heatmap:
if tracks[0].boxes.id is None:
if self.view_img and self.env_check:
self.display_frames()
return
else:
return
return
self.heatmap *= self.decay_factor # decay factor
self.extract_results(tracks)
self.annotator = Annotator(self.im0, self.count_txt_thickness, None)
@ -207,23 +204,21 @@ class Heatmap:
# Count objects
if len(self.count_reg_pts) == 4:
if self.counting_region.contains(Point(track_line[-1])):
if track_id not in self.counting_list:
self.counting_list.append(track_id)
if box[0] < self.counting_region.centroid.x:
self.out_counts += 1
else:
self.in_counts += 1
if self.counting_region.contains(Point(track_line[-1])) and track_id not in self.counting_list:
self.counting_list.append(track_id)
if box[0] < self.counting_region.centroid.x:
self.out_counts += 1
else:
self.in_counts += 1
elif len(self.count_reg_pts) == 2:
distance = Point(track_line[-1]).distance(self.counting_region)
if distance < self.line_dist_thresh:
if track_id not in self.counting_list:
self.counting_list.append(track_id)
if box[0] < self.counting_region.centroid.x:
self.out_counts += 1
else:
self.in_counts += 1
if distance < self.line_dist_thresh and track_id not in self.counting_list:
self.counting_list.append(track_id)
if box[0] < self.counting_region.centroid.x:
self.out_counts += 1
else:
self.in_counts += 1
else:
for box, cls in zip(self.boxes, self.clss):
if self.shape == "circle":
@ -244,8 +239,8 @@ class Heatmap:
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 = f"In Count : {self.in_counts}"
outcount_label = f"OutCount : {self.out_counts}"
# Display counts based on user choice
counts_label = None
@ -256,7 +251,7 @@ class Heatmap:
elif not self.view_out_counts:
counts_label = incount_label
else:
counts_label = incount_label + " " + outcount_label
counts_label = f"{incount_label} {outcount_label}"
if self.count_reg_pts is not None and counts_label is not None:
self.annotator.count_labels(

59
ultralytics/solutions/object_counter.py
View file

@ -139,11 +139,14 @@ class ObjectCounter:
# global is_drawing, selected_point
if event == cv2.EVENT_LBUTTONDOWN:
for i, point in enumerate(self.reg_pts):
if isinstance(point, (tuple, list)) and len(point) >= 2:
if abs(x - point[0]) < 10 and abs(y - point[1]) < 10:
self.selected_point = i
self.is_drawing = True
break
if (
isinstance(point, (tuple, list))
and len(point) >= 2
and (abs(x - point[0]) < 10 and abs(y - point[1]) < 10)
):
self.selected_point = i
self.is_drawing = True
break
elif event == cv2.EVENT_MOUSEMOVE:
if self.is_drawing and self.selected_point is not None:
@ -166,9 +169,8 @@ 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
# Draw bounding box
self.annotator.box_label(box, label=f"{track_id}:{self.names[cls]}", color=colors(int(cls), True))
# Draw Tracks
track_line = self.track_history[track_id]
@ -186,28 +188,29 @@ class ObjectCounter:
# Count objects
if len(self.reg_pts) == 4:
if prev_position is not None:
if self.counting_region.contains(Point(track_line[-1])):
if track_id not in self.counting_list:
self.counting_list.append(track_id)
if (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0:
self.in_counts += 1
else:
self.out_counts += 1
if (
prev_position is not None
and self.counting_region.contains(Point(track_line[-1]))
and track_id not in self.counting_list
):
self.counting_list.append(track_id)
if (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0:
self.in_counts += 1
else:
self.out_counts += 1
elif len(self.reg_pts) == 2:
if prev_position is not None:
distance = Point(track_line[-1]).distance(self.counting_region)
if distance < self.line_dist_thresh:
if track_id not in self.counting_list:
self.counting_list.append(track_id)
if (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0:
self.in_counts += 1
else:
self.out_counts += 1
if distance < self.line_dist_thresh and track_id not in self.counting_list:
self.counting_list.append(track_id)
if (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0:
self.in_counts += 1
else:
self.out_counts += 1
incount_label = "In Count : " + f"{self.in_counts}"
outcount_label = "OutCount : " + f"{self.out_counts}"
incount_label = f"In Count : {self.in_counts}"
outcount_label = f"OutCount : {self.out_counts}"
# Display counts based on user choice
counts_label = None
@ -218,7 +221,7 @@ class ObjectCounter:
elif not self.view_out_counts:
counts_label = incount_label
else:
counts_label = incount_label + " " + outcount_label
counts_label = f"{incount_label} {outcount_label}"
if counts_label is not None:
self.annotator.count_labels(
@ -254,9 +257,7 @@ class ObjectCounter:
if tracks[0].boxes.id is None:
if self.view_img:
self.display_frames()
return
else:
return
return
self.extract_and_process_tracks(tracks)
if self.view_img:

43
ultralytics/solutions/speed_estimation.py
View file

@ -114,9 +114,7 @@ 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 = f"{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,28 +130,28 @@ class SpeedEstimator:
track (list): tracking history for tracks path drawing
"""
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 not self.reg_pts[0][0] < track[-1][0] < self.reg_pts[1][0]:
return
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"
else:
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)
if self.trk_previous_times[trk_id] != 0 and direction != "unknown" and trk_id not in self.trk_idslist:
self.trk_idslist.append(trk_id)
time_difference = time() - self.trk_previous_times[trk_id]
if time_difference > 0:
dist_difference = np.abs(track[-1][1] - self.trk_previous_points[trk_id][1])
speed = dist_difference / time_difference
self.dist_data[trk_id] = speed
time_difference = time() - self.trk_previous_times[trk_id]
if time_difference > 0:
dist_difference = np.abs(track[-1][1] - self.trk_previous_points[trk_id][1])
speed = dist_difference / time_difference
self.dist_data[trk_id] = speed
self.trk_previous_times[trk_id] = time()
self.trk_previous_points[trk_id] = track[-1]
self.trk_previous_times[trk_id] = time()
self.trk_previous_points[trk_id] = track[-1]
def estimate_speed(self, im0, tracks):
"""
@ -166,10 +164,7 @@ class SpeedEstimator:
if tracks[0].boxes.id is None:
if self.view_img and self.env_check:
self.display_frames()
return
else:
return
return
self.extract_tracks(tracks)
self.annotator = Annotator(self.im0, line_width=2)