ultralytics 8.0.237__py3-none-any.whl → 8.0.239__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Potentially problematic release.
This version of ultralytics might be problematic. Click here for more details.
- ultralytics/__init__.py +2 -2
- ultralytics/cfg/__init__.py +241 -138
- ultralytics/cfg/datasets/DOTAv1.5.yaml +1 -1
- ultralytics/cfg/datasets/DOTAv1.yaml +1 -1
- ultralytics/cfg/datasets/dota8.yaml +34 -0
- ultralytics/data/__init__.py +9 -2
- ultralytics/data/annotator.py +4 -4
- ultralytics/data/augment.py +186 -169
- ultralytics/data/base.py +54 -48
- ultralytics/data/build.py +34 -23
- ultralytics/data/converter.py +242 -70
- ultralytics/data/dataset.py +117 -95
- ultralytics/data/explorer/__init__.py +5 -0
- ultralytics/data/explorer/explorer.py +170 -97
- ultralytics/data/explorer/gui/__init__.py +1 -0
- ultralytics/data/explorer/gui/dash.py +146 -76
- ultralytics/data/explorer/utils.py +87 -25
- ultralytics/data/loaders.py +75 -62
- ultralytics/data/split_dota.py +44 -36
- ultralytics/data/utils.py +160 -142
- ultralytics/engine/exporter.py +348 -292
- ultralytics/engine/model.py +102 -66
- ultralytics/engine/predictor.py +74 -55
- ultralytics/engine/results.py +63 -40
- ultralytics/engine/trainer.py +192 -144
- ultralytics/engine/tuner.py +66 -59
- ultralytics/engine/validator.py +31 -26
- ultralytics/hub/__init__.py +54 -31
- ultralytics/hub/auth.py +28 -25
- ultralytics/hub/session.py +282 -133
- ultralytics/hub/utils.py +64 -42
- ultralytics/models/__init__.py +1 -1
- ultralytics/models/fastsam/__init__.py +1 -1
- ultralytics/models/fastsam/model.py +6 -6
- ultralytics/models/fastsam/predict.py +3 -2
- ultralytics/models/fastsam/prompt.py +55 -48
- ultralytics/models/fastsam/val.py +1 -1
- ultralytics/models/nas/__init__.py +1 -1
- ultralytics/models/nas/model.py +9 -8
- ultralytics/models/nas/predict.py +8 -6
- ultralytics/models/nas/val.py +11 -9
- ultralytics/models/rtdetr/__init__.py +1 -1
- ultralytics/models/rtdetr/model.py +11 -9
- ultralytics/models/rtdetr/train.py +18 -16
- ultralytics/models/rtdetr/val.py +25 -19
- ultralytics/models/sam/__init__.py +1 -1
- ultralytics/models/sam/amg.py +13 -14
- ultralytics/models/sam/build.py +44 -42
- ultralytics/models/sam/model.py +6 -6
- ultralytics/models/sam/modules/decoders.py +6 -4
- ultralytics/models/sam/modules/encoders.py +37 -35
- ultralytics/models/sam/modules/sam.py +5 -4
- ultralytics/models/sam/modules/tiny_encoder.py +95 -73
- ultralytics/models/sam/modules/transformer.py +3 -2
- ultralytics/models/sam/predict.py +39 -27
- ultralytics/models/utils/loss.py +99 -95
- ultralytics/models/utils/ops.py +34 -31
- ultralytics/models/yolo/__init__.py +1 -1
- ultralytics/models/yolo/classify/__init__.py +1 -1
- ultralytics/models/yolo/classify/predict.py +8 -6
- ultralytics/models/yolo/classify/train.py +37 -31
- ultralytics/models/yolo/classify/val.py +26 -24
- ultralytics/models/yolo/detect/__init__.py +1 -1
- ultralytics/models/yolo/detect/predict.py +8 -6
- ultralytics/models/yolo/detect/train.py +47 -37
- ultralytics/models/yolo/detect/val.py +100 -82
- ultralytics/models/yolo/model.py +31 -25
- ultralytics/models/yolo/obb/__init__.py +1 -1
- ultralytics/models/yolo/obb/predict.py +13 -12
- ultralytics/models/yolo/obb/train.py +3 -3
- ultralytics/models/yolo/obb/val.py +80 -58
- ultralytics/models/yolo/pose/__init__.py +1 -1
- ultralytics/models/yolo/pose/predict.py +17 -12
- ultralytics/models/yolo/pose/train.py +28 -25
- ultralytics/models/yolo/pose/val.py +91 -64
- ultralytics/models/yolo/segment/__init__.py +1 -1
- ultralytics/models/yolo/segment/predict.py +10 -8
- ultralytics/models/yolo/segment/train.py +16 -15
- ultralytics/models/yolo/segment/val.py +90 -68
- ultralytics/nn/__init__.py +26 -6
- ultralytics/nn/autobackend.py +144 -112
- ultralytics/nn/modules/__init__.py +96 -13
- ultralytics/nn/modules/block.py +28 -7
- ultralytics/nn/modules/conv.py +41 -23
- ultralytics/nn/modules/head.py +67 -59
- ultralytics/nn/modules/transformer.py +49 -32
- ultralytics/nn/modules/utils.py +20 -15
- ultralytics/nn/tasks.py +215 -141
- ultralytics/solutions/ai_gym.py +59 -47
- ultralytics/solutions/distance_calculation.py +22 -15
- ultralytics/solutions/heatmap.py +76 -54
- ultralytics/solutions/object_counter.py +46 -39
- ultralytics/solutions/speed_estimation.py +13 -16
- ultralytics/trackers/__init__.py +1 -1
- ultralytics/trackers/basetrack.py +1 -0
- ultralytics/trackers/bot_sort.py +2 -1
- ultralytics/trackers/byte_tracker.py +10 -7
- ultralytics/trackers/track.py +7 -7
- ultralytics/trackers/utils/gmc.py +25 -25
- ultralytics/trackers/utils/kalman_filter.py +85 -42
- ultralytics/trackers/utils/matching.py +8 -7
- ultralytics/utils/__init__.py +173 -151
- ultralytics/utils/autobatch.py +10 -10
- ultralytics/utils/benchmarks.py +76 -86
- ultralytics/utils/callbacks/__init__.py +1 -1
- ultralytics/utils/callbacks/base.py +29 -29
- ultralytics/utils/callbacks/clearml.py +51 -43
- ultralytics/utils/callbacks/comet.py +81 -66
- ultralytics/utils/callbacks/dvc.py +33 -26
- ultralytics/utils/callbacks/hub.py +44 -26
- ultralytics/utils/callbacks/mlflow.py +31 -24
- ultralytics/utils/callbacks/neptune.py +35 -25
- ultralytics/utils/callbacks/raytune.py +9 -4
- ultralytics/utils/callbacks/tensorboard.py +16 -11
- ultralytics/utils/callbacks/wb.py +39 -33
- ultralytics/utils/checks.py +189 -141
- ultralytics/utils/dist.py +15 -12
- ultralytics/utils/downloads.py +112 -96
- ultralytics/utils/errors.py +1 -1
- ultralytics/utils/files.py +11 -11
- ultralytics/utils/instance.py +22 -22
- ultralytics/utils/loss.py +117 -67
- ultralytics/utils/metrics.py +224 -158
- ultralytics/utils/ops.py +39 -29
- ultralytics/utils/patches.py +3 -3
- ultralytics/utils/plotting.py +217 -120
- ultralytics/utils/tal.py +19 -13
- ultralytics/utils/torch_utils.py +138 -109
- ultralytics/utils/triton.py +12 -10
- ultralytics/utils/tuner.py +49 -47
- {ultralytics-8.0.237.dist-info → ultralytics-8.0.239.dist-info}/METADATA +5 -4
- ultralytics-8.0.239.dist-info/RECORD +188 -0
- ultralytics-8.0.237.dist-info/RECORD +0 -187
- {ultralytics-8.0.237.dist-info → ultralytics-8.0.239.dist-info}/LICENSE +0 -0
- {ultralytics-8.0.237.dist-info → ultralytics-8.0.239.dist-info}/WHEEL +0 -0
- {ultralytics-8.0.237.dist-info → ultralytics-8.0.239.dist-info}/entry_points.txt +0 -0
- {ultralytics-8.0.237.dist-info → ultralytics-8.0.239.dist-info}/top_level.txt +0 -0
ultralytics/solutions/ai_gym.py
CHANGED
|
@@ -26,7 +26,7 @@ class AIGym:
|
|
|
26
26
|
self.angle = None
|
|
27
27
|
self.count = None
|
|
28
28
|
self.stage = None
|
|
29
|
-
self.pose_type =
|
|
29
|
+
self.pose_type = "pushup"
|
|
30
30
|
self.kpts_to_check = None
|
|
31
31
|
|
|
32
32
|
# Visual Information
|
|
@@ -36,13 +36,15 @@ class AIGym:
|
|
|
36
36
|
# Check if environment support imshow
|
|
37
37
|
self.env_check = check_imshow(warn=True)
|
|
38
38
|
|
|
39
|
-
def set_args(
|
|
40
|
-
|
|
41
|
-
|
|
42
|
-
|
|
43
|
-
|
|
44
|
-
|
|
45
|
-
|
|
39
|
+
def set_args(
|
|
40
|
+
self,
|
|
41
|
+
kpts_to_check,
|
|
42
|
+
line_thickness=2,
|
|
43
|
+
view_img=False,
|
|
44
|
+
pose_up_angle=145.0,
|
|
45
|
+
pose_down_angle=90.0,
|
|
46
|
+
pose_type="pullup",
|
|
47
|
+
):
|
|
46
48
|
"""
|
|
47
49
|
Configures the AIGym line_thickness, save image and view image parameters
|
|
48
50
|
Args:
|
|
@@ -72,65 +74,75 @@ class AIGym:
|
|
|
72
74
|
if frame_count == 1:
|
|
73
75
|
self.count = [0] * len(results[0])
|
|
74
76
|
self.angle = [0] * len(results[0])
|
|
75
|
-
self.stage = [
|
|
77
|
+
self.stage = ["-" for _ in results[0]]
|
|
76
78
|
self.keypoints = results[0].keypoints.data
|
|
77
79
|
self.annotator = Annotator(im0, line_width=2)
|
|
78
80
|
|
|
79
81
|
for ind, k in enumerate(reversed(self.keypoints)):
|
|
80
|
-
if self.pose_type ==
|
|
81
|
-
self.angle[ind] = self.annotator.estimate_pose_angle(
|
|
82
|
-
|
|
83
|
-
|
|
82
|
+
if self.pose_type == "pushup" or self.pose_type == "pullup":
|
|
83
|
+
self.angle[ind] = self.annotator.estimate_pose_angle(
|
|
84
|
+
k[int(self.kpts_to_check[0])].cpu(),
|
|
85
|
+
k[int(self.kpts_to_check[1])].cpu(),
|
|
86
|
+
k[int(self.kpts_to_check[2])].cpu(),
|
|
87
|
+
)
|
|
84
88
|
self.im0 = self.annotator.draw_specific_points(k, self.kpts_to_check, shape=(640, 640), radius=10)
|
|
85
89
|
|
|
86
|
-
if self.pose_type ==
|
|
87
|
-
self.angle[ind] = self.annotator.estimate_pose_angle(
|
|
88
|
-
|
|
89
|
-
|
|
90
|
+
if self.pose_type == "abworkout":
|
|
91
|
+
self.angle[ind] = self.annotator.estimate_pose_angle(
|
|
92
|
+
k[int(self.kpts_to_check[0])].cpu(),
|
|
93
|
+
k[int(self.kpts_to_check[1])].cpu(),
|
|
94
|
+
k[int(self.kpts_to_check[2])].cpu(),
|
|
95
|
+
)
|
|
90
96
|
self.im0 = self.annotator.draw_specific_points(k, self.kpts_to_check, shape=(640, 640), radius=10)
|
|
91
97
|
if self.angle[ind] > self.poseup_angle:
|
|
92
|
-
self.stage[ind] =
|
|
93
|
-
if self.angle[ind] < self.posedown_angle and self.stage[ind] ==
|
|
94
|
-
self.stage[ind] =
|
|
98
|
+
self.stage[ind] = "down"
|
|
99
|
+
if self.angle[ind] < self.posedown_angle and self.stage[ind] == "down":
|
|
100
|
+
self.stage[ind] = "up"
|
|
95
101
|
self.count[ind] += 1
|
|
96
|
-
self.annotator.plot_angle_and_count_and_stage(
|
|
97
|
-
|
|
98
|
-
|
|
99
|
-
|
|
100
|
-
|
|
101
|
-
|
|
102
|
-
|
|
102
|
+
self.annotator.plot_angle_and_count_and_stage(
|
|
103
|
+
angle_text=self.angle[ind],
|
|
104
|
+
count_text=self.count[ind],
|
|
105
|
+
stage_text=self.stage[ind],
|
|
106
|
+
center_kpt=k[int(self.kpts_to_check[1])],
|
|
107
|
+
line_thickness=self.tf,
|
|
108
|
+
)
|
|
109
|
+
|
|
110
|
+
if self.pose_type == "pushup":
|
|
103
111
|
if self.angle[ind] > self.poseup_angle:
|
|
104
|
-
self.stage[ind] =
|
|
105
|
-
if self.angle[ind] < self.posedown_angle and self.stage[ind] ==
|
|
106
|
-
self.stage[ind] =
|
|
112
|
+
self.stage[ind] = "up"
|
|
113
|
+
if self.angle[ind] < self.posedown_angle and self.stage[ind] == "up":
|
|
114
|
+
self.stage[ind] = "down"
|
|
107
115
|
self.count[ind] += 1
|
|
108
|
-
self.annotator.plot_angle_and_count_and_stage(
|
|
109
|
-
|
|
110
|
-
|
|
111
|
-
|
|
112
|
-
|
|
113
|
-
|
|
116
|
+
self.annotator.plot_angle_and_count_and_stage(
|
|
117
|
+
angle_text=self.angle[ind],
|
|
118
|
+
count_text=self.count[ind],
|
|
119
|
+
stage_text=self.stage[ind],
|
|
120
|
+
center_kpt=k[int(self.kpts_to_check[1])],
|
|
121
|
+
line_thickness=self.tf,
|
|
122
|
+
)
|
|
123
|
+
if self.pose_type == "pullup":
|
|
114
124
|
if self.angle[ind] > self.poseup_angle:
|
|
115
|
-
self.stage[ind] =
|
|
116
|
-
if self.angle[ind] < self.posedown_angle and self.stage[ind] ==
|
|
117
|
-
self.stage[ind] =
|
|
125
|
+
self.stage[ind] = "down"
|
|
126
|
+
if self.angle[ind] < self.posedown_angle and self.stage[ind] == "down":
|
|
127
|
+
self.stage[ind] = "up"
|
|
118
128
|
self.count[ind] += 1
|
|
119
|
-
self.annotator.plot_angle_and_count_and_stage(
|
|
120
|
-
|
|
121
|
-
|
|
122
|
-
|
|
123
|
-
|
|
129
|
+
self.annotator.plot_angle_and_count_and_stage(
|
|
130
|
+
angle_text=self.angle[ind],
|
|
131
|
+
count_text=self.count[ind],
|
|
132
|
+
stage_text=self.stage[ind],
|
|
133
|
+
center_kpt=k[int(self.kpts_to_check[1])],
|
|
134
|
+
line_thickness=self.tf,
|
|
135
|
+
)
|
|
124
136
|
|
|
125
137
|
self.annotator.kpts(k, shape=(640, 640), radius=1, kpt_line=True)
|
|
126
138
|
|
|
127
139
|
if self.env_check and self.view_img:
|
|
128
|
-
cv2.imshow(
|
|
129
|
-
if cv2.waitKey(1) & 0xFF == ord(
|
|
140
|
+
cv2.imshow("Ultralytics YOLOv8 AI GYM", self.im0)
|
|
141
|
+
if cv2.waitKey(1) & 0xFF == ord("q"):
|
|
130
142
|
return
|
|
131
143
|
|
|
132
144
|
return self.im0
|
|
133
145
|
|
|
134
146
|
|
|
135
|
-
if __name__ ==
|
|
147
|
+
if __name__ == "__main__":
|
|
136
148
|
AIGym()
|
|
@@ -4,6 +4,7 @@ import math
|
|
|
4
4
|
|
|
5
5
|
import cv2
|
|
6
6
|
|
|
7
|
+
from ultralytics.utils.checks import check_imshow
|
|
7
8
|
from ultralytics.utils.plotting import Annotator, colors
|
|
8
9
|
|
|
9
10
|
|
|
@@ -37,13 +38,18 @@ class DistanceCalculation:
|
|
|
37
38
|
self.left_mouse_count = 0
|
|
38
39
|
self.selected_boxes = {}
|
|
39
40
|
|
|
40
|
-
|
|
41
|
-
|
|
42
|
-
|
|
43
|
-
|
|
44
|
-
|
|
45
|
-
|
|
46
|
-
|
|
41
|
+
# Check if environment support imshow
|
|
42
|
+
self.env_check = check_imshow(warn=True)
|
|
43
|
+
|
|
44
|
+
def set_args(
|
|
45
|
+
self,
|
|
46
|
+
names,
|
|
47
|
+
pixels_per_meter=10,
|
|
48
|
+
view_img=False,
|
|
49
|
+
line_thickness=2,
|
|
50
|
+
line_color=(255, 255, 0),
|
|
51
|
+
centroid_color=(255, 0, 255),
|
|
52
|
+
):
|
|
47
53
|
"""
|
|
48
54
|
Configures the distance calculation and display parameters.
|
|
49
55
|
|
|
@@ -125,8 +131,9 @@ class DistanceCalculation:
|
|
|
125
131
|
distance (float): Distance between two centroids
|
|
126
132
|
"""
|
|
127
133
|
cv2.rectangle(self.im0, (15, 25), (280, 70), (255, 255, 255), -1)
|
|
128
|
-
cv2.putText(
|
|
129
|
-
|
|
134
|
+
cv2.putText(
|
|
135
|
+
self.im0, f"Distance : {distance:.2f}m", (20, 55), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 0), 2, cv2.LINE_AA
|
|
136
|
+
)
|
|
130
137
|
cv2.line(self.im0, self.centroids[0], self.centroids[1], self.line_color, 3)
|
|
131
138
|
cv2.circle(self.im0, self.centroids[0], 6, self.centroid_color, -1)
|
|
132
139
|
cv2.circle(self.im0, self.centroids[1], 6, self.centroid_color, -1)
|
|
@@ -168,20 +175,20 @@ class DistanceCalculation:
|
|
|
168
175
|
|
|
169
176
|
self.centroids = []
|
|
170
177
|
|
|
171
|
-
if self.view_img:
|
|
178
|
+
if self.view_img and self.env_check:
|
|
172
179
|
self.display_frames()
|
|
173
180
|
|
|
174
181
|
return im0
|
|
175
182
|
|
|
176
183
|
def display_frames(self):
|
|
177
184
|
"""Display frame."""
|
|
178
|
-
cv2.namedWindow(
|
|
179
|
-
cv2.setMouseCallback(
|
|
180
|
-
cv2.imshow(
|
|
185
|
+
cv2.namedWindow("Ultralytics Distance Estimation")
|
|
186
|
+
cv2.setMouseCallback("Ultralytics Distance Estimation", self.mouse_event_for_distance)
|
|
187
|
+
cv2.imshow("Ultralytics Distance Estimation", self.im0)
|
|
181
188
|
|
|
182
|
-
if cv2.waitKey(1) & 0xFF == ord(
|
|
189
|
+
if cv2.waitKey(1) & 0xFF == ord("q"):
|
|
183
190
|
return
|
|
184
191
|
|
|
185
192
|
|
|
186
|
-
if __name__ ==
|
|
193
|
+
if __name__ == "__main__":
|
|
187
194
|
DistanceCalculation()
|
ultralytics/solutions/heatmap.py
CHANGED
|
@@ -8,7 +8,7 @@ import numpy as np
|
|
|
8
8
|
from ultralytics.utils.checks import check_imshow, check_requirements
|
|
9
9
|
from ultralytics.utils.plotting import Annotator
|
|
10
10
|
|
|
11
|
-
check_requirements(
|
|
11
|
+
check_requirements("shapely>=2.0.0")
|
|
12
12
|
|
|
13
13
|
from shapely.geometry import LineString, Point, Polygon
|
|
14
14
|
|
|
@@ -22,12 +22,14 @@ class Heatmap:
|
|
|
22
22
|
# Visual information
|
|
23
23
|
self.annotator = None
|
|
24
24
|
self.view_img = False
|
|
25
|
-
self.shape =
|
|
25
|
+
self.shape = "circle"
|
|
26
26
|
|
|
27
27
|
# Image information
|
|
28
28
|
self.imw = None
|
|
29
29
|
self.imh = None
|
|
30
30
|
self.im0 = None
|
|
31
|
+
self.view_in_counts = True
|
|
32
|
+
self.view_out_counts = True
|
|
31
33
|
|
|
32
34
|
# Heatmap colormap and heatmap np array
|
|
33
35
|
self.colormap = None
|
|
@@ -61,21 +63,25 @@ class Heatmap:
|
|
|
61
63
|
# Check if environment support imshow
|
|
62
64
|
self.env_check = check_imshow(warn=True)
|
|
63
65
|
|
|
64
|
-
def set_args(
|
|
65
|
-
|
|
66
|
-
|
|
67
|
-
|
|
68
|
-
|
|
69
|
-
|
|
70
|
-
|
|
71
|
-
|
|
72
|
-
|
|
73
|
-
|
|
74
|
-
|
|
75
|
-
|
|
76
|
-
|
|
77
|
-
|
|
78
|
-
|
|
66
|
+
def set_args(
|
|
67
|
+
self,
|
|
68
|
+
imw,
|
|
69
|
+
imh,
|
|
70
|
+
colormap=cv2.COLORMAP_JET,
|
|
71
|
+
heatmap_alpha=0.5,
|
|
72
|
+
view_img=False,
|
|
73
|
+
view_in_counts=True,
|
|
74
|
+
view_out_counts=True,
|
|
75
|
+
count_reg_pts=None,
|
|
76
|
+
count_txt_thickness=2,
|
|
77
|
+
count_txt_color=(0, 0, 0),
|
|
78
|
+
count_color=(255, 255, 255),
|
|
79
|
+
count_reg_color=(255, 0, 255),
|
|
80
|
+
region_thickness=5,
|
|
81
|
+
line_dist_thresh=15,
|
|
82
|
+
decay_factor=0.99,
|
|
83
|
+
shape="circle",
|
|
84
|
+
):
|
|
79
85
|
"""
|
|
80
86
|
Configures the heatmap colormap, width, height and display parameters.
|
|
81
87
|
|
|
@@ -85,6 +91,8 @@ class Heatmap:
|
|
|
85
91
|
imh (int): The height of the frame.
|
|
86
92
|
heatmap_alpha (float): alpha value for heatmap display
|
|
87
93
|
view_img (bool): Flag indicating frame display
|
|
94
|
+
view_in_counts (bool): Flag to control whether to display the incounts on video stream.
|
|
95
|
+
view_out_counts (bool): Flag to control whether to display the outcounts on video stream.
|
|
88
96
|
count_reg_pts (list): Object counting region points
|
|
89
97
|
count_txt_thickness (int): Text thickness for object counting display
|
|
90
98
|
count_txt_color (RGB color): count text color value
|
|
@@ -99,24 +107,25 @@ class Heatmap:
|
|
|
99
107
|
self.imh = imh
|
|
100
108
|
self.heatmap_alpha = heatmap_alpha
|
|
101
109
|
self.view_img = view_img
|
|
110
|
+
self.view_in_counts = view_in_counts
|
|
111
|
+
self.view_out_counts = view_out_counts
|
|
102
112
|
self.colormap = colormap
|
|
103
113
|
|
|
104
114
|
# Region and line selection
|
|
105
115
|
if count_reg_pts is not None:
|
|
106
|
-
|
|
107
116
|
if len(count_reg_pts) == 2:
|
|
108
|
-
print(
|
|
117
|
+
print("Line Counter Initiated.")
|
|
109
118
|
self.count_reg_pts = count_reg_pts
|
|
110
119
|
self.counting_region = LineString(count_reg_pts)
|
|
111
120
|
|
|
112
121
|
elif len(count_reg_pts) == 4:
|
|
113
|
-
print(
|
|
122
|
+
print("Region Counter Initiated.")
|
|
114
123
|
self.count_reg_pts = count_reg_pts
|
|
115
124
|
self.counting_region = Polygon(self.count_reg_pts)
|
|
116
125
|
|
|
117
126
|
else:
|
|
118
|
-
print(
|
|
119
|
-
print(
|
|
127
|
+
print("Region or line points Invalid, 2 or 4 points supported")
|
|
128
|
+
print("Using Line Counter Now")
|
|
120
129
|
self.counting_region = Polygon([(20, 400), (1260, 400)]) # dummy points
|
|
121
130
|
|
|
122
131
|
# Heatmap new frame
|
|
@@ -132,10 +141,10 @@ class Heatmap:
|
|
|
132
141
|
self.shape = shape
|
|
133
142
|
|
|
134
143
|
# shape of heatmap, if not selected
|
|
135
|
-
if self.shape not in [
|
|
144
|
+
if self.shape not in ["circle", "rect"]:
|
|
136
145
|
print("Unknown shape value provided, 'circle' & 'rect' supported")
|
|
137
|
-
print(
|
|
138
|
-
self.shape =
|
|
146
|
+
print("Using Circular shape now")
|
|
147
|
+
self.shape = "circle"
|
|
139
148
|
|
|
140
149
|
def extract_results(self, tracks):
|
|
141
150
|
"""
|
|
@@ -169,26 +178,26 @@ class Heatmap:
|
|
|
169
178
|
self.annotator = Annotator(self.im0, self.count_txt_thickness, None)
|
|
170
179
|
|
|
171
180
|
if self.count_reg_pts is not None:
|
|
172
|
-
|
|
173
181
|
# Draw counting region
|
|
174
|
-
self.
|
|
175
|
-
|
|
176
|
-
|
|
182
|
+
if self.view_in_counts or self.view_out_counts:
|
|
183
|
+
self.annotator.draw_region(
|
|
184
|
+
reg_pts=self.count_reg_pts, color=self.region_color, thickness=self.region_thickness
|
|
185
|
+
)
|
|
177
186
|
|
|
178
187
|
for box, cls, track_id in zip(self.boxes, self.clss, self.track_ids):
|
|
179
|
-
|
|
180
|
-
if self.shape == 'circle':
|
|
188
|
+
if self.shape == "circle":
|
|
181
189
|
center = (int((box[0] + box[2]) // 2), int((box[1] + box[3]) // 2))
|
|
182
190
|
radius = min(int(box[2]) - int(box[0]), int(box[3]) - int(box[1])) // 2
|
|
183
191
|
|
|
184
|
-
y, x = np.ogrid[0:self.heatmap.shape[0], 0:self.heatmap.shape[1]]
|
|
185
|
-
mask = (x - center[0]) ** 2 + (y - center[1]) ** 2 <= radius
|
|
192
|
+
y, x = np.ogrid[0 : self.heatmap.shape[0], 0 : self.heatmap.shape[1]]
|
|
193
|
+
mask = (x - center[0]) ** 2 + (y - center[1]) ** 2 <= radius**2
|
|
186
194
|
|
|
187
|
-
self.heatmap[int(box[1]):int(box[3]), int(box[0]):int(box[2])] +=
|
|
188
|
-
|
|
195
|
+
self.heatmap[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])] += (
|
|
196
|
+
2 * mask[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])]
|
|
197
|
+
)
|
|
189
198
|
|
|
190
199
|
else:
|
|
191
|
-
self.heatmap[int(box[1]):int(box[3]), int(box[0]):int(box[2])] += 2
|
|
200
|
+
self.heatmap[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])] += 2
|
|
192
201
|
|
|
193
202
|
# Store tracking hist
|
|
194
203
|
track_line = self.track_history[track_id]
|
|
@@ -217,32 +226,45 @@ class Heatmap:
|
|
|
217
226
|
self.in_counts += 1
|
|
218
227
|
else:
|
|
219
228
|
for box, cls in zip(self.boxes, self.clss):
|
|
220
|
-
|
|
221
|
-
if self.shape == 'circle':
|
|
229
|
+
if self.shape == "circle":
|
|
222
230
|
center = (int((box[0] + box[2]) // 2), int((box[1] + box[3]) // 2))
|
|
223
231
|
radius = min(int(box[2]) - int(box[0]), int(box[3]) - int(box[1])) // 2
|
|
224
232
|
|
|
225
|
-
y, x = np.ogrid[0:self.heatmap.shape[0], 0:self.heatmap.shape[1]]
|
|
226
|
-
mask = (x - center[0]) ** 2 + (y - center[1]) ** 2 <= radius
|
|
233
|
+
y, x = np.ogrid[0 : self.heatmap.shape[0], 0 : self.heatmap.shape[1]]
|
|
234
|
+
mask = (x - center[0]) ** 2 + (y - center[1]) ** 2 <= radius**2
|
|
227
235
|
|
|
228
|
-
self.heatmap[int(box[1]):int(box[3]), int(box[0]):int(box[2])] +=
|
|
229
|
-
|
|
236
|
+
self.heatmap[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])] += (
|
|
237
|
+
2 * mask[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])]
|
|
238
|
+
)
|
|
230
239
|
|
|
231
240
|
else:
|
|
232
|
-
self.heatmap[int(box[1]):int(box[3]), int(box[0]):int(box[2])] += 2
|
|
241
|
+
self.heatmap[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])] += 2
|
|
233
242
|
|
|
234
243
|
# Normalize, apply colormap to heatmap and combine with original image
|
|
235
244
|
heatmap_normalized = cv2.normalize(self.heatmap, None, 0, 255, cv2.NORM_MINMAX)
|
|
236
245
|
heatmap_colored = cv2.applyColorMap(heatmap_normalized.astype(np.uint8), self.colormap)
|
|
237
246
|
|
|
238
|
-
|
|
239
|
-
|
|
240
|
-
|
|
241
|
-
|
|
242
|
-
|
|
243
|
-
|
|
244
|
-
|
|
245
|
-
|
|
247
|
+
incount_label = "In Count : " + f"{self.in_counts}"
|
|
248
|
+
outcount_label = "OutCount : " + f"{self.out_counts}"
|
|
249
|
+
|
|
250
|
+
# Display counts based on user choice
|
|
251
|
+
counts_label = None
|
|
252
|
+
if not self.view_in_counts and not self.view_out_counts:
|
|
253
|
+
counts_label = None
|
|
254
|
+
elif not self.view_in_counts:
|
|
255
|
+
counts_label = outcount_label
|
|
256
|
+
elif not self.view_out_counts:
|
|
257
|
+
counts_label = incount_label
|
|
258
|
+
else:
|
|
259
|
+
counts_label = incount_label + " " + outcount_label
|
|
260
|
+
|
|
261
|
+
if self.count_reg_pts is not None and counts_label is not None:
|
|
262
|
+
self.annotator.count_labels(
|
|
263
|
+
counts=counts_label,
|
|
264
|
+
count_txt_size=self.count_txt_thickness,
|
|
265
|
+
txt_color=self.count_txt_color,
|
|
266
|
+
color=self.count_color,
|
|
267
|
+
)
|
|
246
268
|
|
|
247
269
|
self.im0 = cv2.addWeighted(self.im0, 1 - self.heatmap_alpha, heatmap_colored, self.heatmap_alpha, 0)
|
|
248
270
|
|
|
@@ -253,11 +275,11 @@ class Heatmap:
|
|
|
253
275
|
|
|
254
276
|
def display_frames(self):
|
|
255
277
|
"""Display frame."""
|
|
256
|
-
cv2.imshow(
|
|
278
|
+
cv2.imshow("Ultralytics Heatmap", self.im0)
|
|
257
279
|
|
|
258
|
-
if cv2.waitKey(1) & 0xFF == ord(
|
|
280
|
+
if cv2.waitKey(1) & 0xFF == ord("q"):
|
|
259
281
|
return
|
|
260
282
|
|
|
261
283
|
|
|
262
|
-
if __name__ ==
|
|
284
|
+
if __name__ == "__main__":
|
|
263
285
|
Heatmap()
|
|
@@ -7,7 +7,7 @@ import cv2
|
|
|
7
7
|
from ultralytics.utils.checks import check_imshow, check_requirements
|
|
8
8
|
from ultralytics.utils.plotting import Annotator, colors
|
|
9
9
|
|
|
10
|
-
check_requirements(
|
|
10
|
+
check_requirements("shapely>=2.0.0")
|
|
11
11
|
|
|
12
12
|
from shapely.geometry import LineString, Point, Polygon
|
|
13
13
|
|
|
@@ -56,22 +56,24 @@ class ObjectCounter:
|
|
|
56
56
|
# Check if environment support imshow
|
|
57
57
|
self.env_check = check_imshow(warn=True)
|
|
58
58
|
|
|
59
|
-
def set_args(
|
|
60
|
-
|
|
61
|
-
|
|
62
|
-
|
|
63
|
-
|
|
64
|
-
|
|
65
|
-
|
|
66
|
-
|
|
67
|
-
|
|
68
|
-
|
|
69
|
-
|
|
70
|
-
|
|
71
|
-
|
|
72
|
-
|
|
73
|
-
|
|
74
|
-
|
|
59
|
+
def set_args(
|
|
60
|
+
self,
|
|
61
|
+
classes_names,
|
|
62
|
+
reg_pts,
|
|
63
|
+
count_reg_color=(255, 0, 255),
|
|
64
|
+
line_thickness=2,
|
|
65
|
+
track_thickness=2,
|
|
66
|
+
view_img=False,
|
|
67
|
+
view_in_counts=True,
|
|
68
|
+
view_out_counts=True,
|
|
69
|
+
draw_tracks=False,
|
|
70
|
+
count_txt_thickness=2,
|
|
71
|
+
count_txt_color=(0, 0, 0),
|
|
72
|
+
count_color=(255, 255, 255),
|
|
73
|
+
track_color=(0, 255, 0),
|
|
74
|
+
region_thickness=5,
|
|
75
|
+
line_dist_thresh=15,
|
|
76
|
+
):
|
|
75
77
|
"""
|
|
76
78
|
Configures the Counter's image, bounding box line thickness, and counting region points.
|
|
77
79
|
|
|
@@ -101,16 +103,16 @@ class ObjectCounter:
|
|
|
101
103
|
|
|
102
104
|
# Region and line selection
|
|
103
105
|
if len(reg_pts) == 2:
|
|
104
|
-
print(
|
|
106
|
+
print("Line Counter Initiated.")
|
|
105
107
|
self.reg_pts = reg_pts
|
|
106
108
|
self.counting_region = LineString(self.reg_pts)
|
|
107
109
|
elif len(reg_pts) == 4:
|
|
108
|
-
print(
|
|
110
|
+
print("Region Counter Initiated.")
|
|
109
111
|
self.reg_pts = reg_pts
|
|
110
112
|
self.counting_region = Polygon(self.reg_pts)
|
|
111
113
|
else:
|
|
112
|
-
print(
|
|
113
|
-
print(
|
|
114
|
+
print("Invalid Region points provided, region_points can be 2 or 4")
|
|
115
|
+
print("Using Line Counter Now")
|
|
114
116
|
self.counting_region = LineString(self.reg_pts)
|
|
115
117
|
|
|
116
118
|
self.names = classes_names
|
|
@@ -153,6 +155,7 @@ class ObjectCounter:
|
|
|
153
155
|
self.selected_point = None
|
|
154
156
|
|
|
155
157
|
def extract_and_process_tracks(self, tracks):
|
|
158
|
+
"""Extracts and processes tracks for object counting in a video stream."""
|
|
156
159
|
boxes = tracks[0].boxes.xyxy.cpu()
|
|
157
160
|
clss = tracks[0].boxes.cls.cpu().tolist()
|
|
158
161
|
track_ids = tracks[0].boxes.id.int().cpu().tolist()
|
|
@@ -163,8 +166,9 @@ class ObjectCounter:
|
|
|
163
166
|
|
|
164
167
|
# Extract tracks
|
|
165
168
|
for box, track_id, cls in zip(boxes, track_ids, clss):
|
|
166
|
-
self.annotator.box_label(
|
|
167
|
-
|
|
169
|
+
self.annotator.box_label(
|
|
170
|
+
box, label=str(track_id) + ":" + self.names[cls], color=colors(int(cls), True)
|
|
171
|
+
) # Draw bounding box
|
|
168
172
|
|
|
169
173
|
# Draw Tracks
|
|
170
174
|
track_line = self.track_history[track_id]
|
|
@@ -174,9 +178,9 @@ class ObjectCounter:
|
|
|
174
178
|
|
|
175
179
|
# Draw track trails
|
|
176
180
|
if self.draw_tracks:
|
|
177
|
-
self.annotator.draw_centroid_and_tracks(
|
|
178
|
-
|
|
179
|
-
|
|
181
|
+
self.annotator.draw_centroid_and_tracks(
|
|
182
|
+
track_line, color=self.track_color, track_thickness=self.track_thickness
|
|
183
|
+
)
|
|
180
184
|
|
|
181
185
|
# Count objects
|
|
182
186
|
if len(self.reg_pts) == 4:
|
|
@@ -198,8 +202,8 @@ class ObjectCounter:
|
|
|
198
202
|
else:
|
|
199
203
|
self.in_counts += 1
|
|
200
204
|
|
|
201
|
-
incount_label =
|
|
202
|
-
outcount_label =
|
|
205
|
+
incount_label = "In Count : " + f"{self.in_counts}"
|
|
206
|
+
outcount_label = "OutCount : " + f"{self.out_counts}"
|
|
203
207
|
|
|
204
208
|
# Display counts based on user choice
|
|
205
209
|
counts_label = None
|
|
@@ -210,24 +214,27 @@ class ObjectCounter:
|
|
|
210
214
|
elif not self.view_out_counts:
|
|
211
215
|
counts_label = incount_label
|
|
212
216
|
else:
|
|
213
|
-
counts_label = incount_label +
|
|
217
|
+
counts_label = incount_label + " " + outcount_label
|
|
214
218
|
|
|
215
219
|
if counts_label is not None:
|
|
216
|
-
self.annotator.count_labels(
|
|
217
|
-
|
|
218
|
-
|
|
219
|
-
|
|
220
|
+
self.annotator.count_labels(
|
|
221
|
+
counts=counts_label,
|
|
222
|
+
count_txt_size=self.count_txt_thickness,
|
|
223
|
+
txt_color=self.count_txt_color,
|
|
224
|
+
color=self.count_color,
|
|
225
|
+
)
|
|
220
226
|
|
|
221
227
|
def display_frames(self):
|
|
222
228
|
"""Display frame."""
|
|
223
229
|
if self.env_check:
|
|
224
|
-
cv2.namedWindow(
|
|
230
|
+
cv2.namedWindow("Ultralytics YOLOv8 Object Counter")
|
|
225
231
|
if len(self.reg_pts) == 4: # only add mouse event If user drawn region
|
|
226
|
-
cv2.setMouseCallback(
|
|
227
|
-
|
|
228
|
-
|
|
232
|
+
cv2.setMouseCallback(
|
|
233
|
+
"Ultralytics YOLOv8 Object Counter", self.mouse_event_for_region, {"region_points": self.reg_pts}
|
|
234
|
+
)
|
|
235
|
+
cv2.imshow("Ultralytics YOLOv8 Object Counter", self.im0)
|
|
229
236
|
# Break Window
|
|
230
|
-
if cv2.waitKey(1) & 0xFF == ord(
|
|
237
|
+
if cv2.waitKey(1) & 0xFF == ord("q"):
|
|
231
238
|
return
|
|
232
239
|
|
|
233
240
|
def start_counting(self, im0, tracks):
|
|
@@ -253,5 +260,5 @@ class ObjectCounter:
|
|
|
253
260
|
return self.im0
|
|
254
261
|
|
|
255
262
|
|
|
256
|
-
if __name__ ==
|
|
263
|
+
if __name__ == "__main__":
|
|
257
264
|
ObjectCounter()
|