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.

Files changed (137) hide show
  1. ultralytics/__init__.py +2 -2
  2. ultralytics/cfg/__init__.py +241 -138
  3. ultralytics/cfg/datasets/DOTAv1.5.yaml +1 -1
  4. ultralytics/cfg/datasets/DOTAv1.yaml +1 -1
  5. ultralytics/cfg/datasets/dota8.yaml +34 -0
  6. ultralytics/data/__init__.py +9 -2
  7. ultralytics/data/annotator.py +4 -4
  8. ultralytics/data/augment.py +186 -169
  9. ultralytics/data/base.py +54 -48
  10. ultralytics/data/build.py +34 -23
  11. ultralytics/data/converter.py +242 -70
  12. ultralytics/data/dataset.py +117 -95
  13. ultralytics/data/explorer/__init__.py +5 -0
  14. ultralytics/data/explorer/explorer.py +170 -97
  15. ultralytics/data/explorer/gui/__init__.py +1 -0
  16. ultralytics/data/explorer/gui/dash.py +146 -76
  17. ultralytics/data/explorer/utils.py +87 -25
  18. ultralytics/data/loaders.py +75 -62
  19. ultralytics/data/split_dota.py +44 -36
  20. ultralytics/data/utils.py +160 -142
  21. ultralytics/engine/exporter.py +348 -292
  22. ultralytics/engine/model.py +102 -66
  23. ultralytics/engine/predictor.py +74 -55
  24. ultralytics/engine/results.py +63 -40
  25. ultralytics/engine/trainer.py +192 -144
  26. ultralytics/engine/tuner.py +66 -59
  27. ultralytics/engine/validator.py +31 -26
  28. ultralytics/hub/__init__.py +54 -31
  29. ultralytics/hub/auth.py +28 -25
  30. ultralytics/hub/session.py +282 -133
  31. ultralytics/hub/utils.py +64 -42
  32. ultralytics/models/__init__.py +1 -1
  33. ultralytics/models/fastsam/__init__.py +1 -1
  34. ultralytics/models/fastsam/model.py +6 -6
  35. ultralytics/models/fastsam/predict.py +3 -2
  36. ultralytics/models/fastsam/prompt.py +55 -48
  37. ultralytics/models/fastsam/val.py +1 -1
  38. ultralytics/models/nas/__init__.py +1 -1
  39. ultralytics/models/nas/model.py +9 -8
  40. ultralytics/models/nas/predict.py +8 -6
  41. ultralytics/models/nas/val.py +11 -9
  42. ultralytics/models/rtdetr/__init__.py +1 -1
  43. ultralytics/models/rtdetr/model.py +11 -9
  44. ultralytics/models/rtdetr/train.py +18 -16
  45. ultralytics/models/rtdetr/val.py +25 -19
  46. ultralytics/models/sam/__init__.py +1 -1
  47. ultralytics/models/sam/amg.py +13 -14
  48. ultralytics/models/sam/build.py +44 -42
  49. ultralytics/models/sam/model.py +6 -6
  50. ultralytics/models/sam/modules/decoders.py +6 -4
  51. ultralytics/models/sam/modules/encoders.py +37 -35
  52. ultralytics/models/sam/modules/sam.py +5 -4
  53. ultralytics/models/sam/modules/tiny_encoder.py +95 -73
  54. ultralytics/models/sam/modules/transformer.py +3 -2
  55. ultralytics/models/sam/predict.py +39 -27
  56. ultralytics/models/utils/loss.py +99 -95
  57. ultralytics/models/utils/ops.py +34 -31
  58. ultralytics/models/yolo/__init__.py +1 -1
  59. ultralytics/models/yolo/classify/__init__.py +1 -1
  60. ultralytics/models/yolo/classify/predict.py +8 -6
  61. ultralytics/models/yolo/classify/train.py +37 -31
  62. ultralytics/models/yolo/classify/val.py +26 -24
  63. ultralytics/models/yolo/detect/__init__.py +1 -1
  64. ultralytics/models/yolo/detect/predict.py +8 -6
  65. ultralytics/models/yolo/detect/train.py +47 -37
  66. ultralytics/models/yolo/detect/val.py +100 -82
  67. ultralytics/models/yolo/model.py +31 -25
  68. ultralytics/models/yolo/obb/__init__.py +1 -1
  69. ultralytics/models/yolo/obb/predict.py +13 -12
  70. ultralytics/models/yolo/obb/train.py +3 -3
  71. ultralytics/models/yolo/obb/val.py +80 -58
  72. ultralytics/models/yolo/pose/__init__.py +1 -1
  73. ultralytics/models/yolo/pose/predict.py +17 -12
  74. ultralytics/models/yolo/pose/train.py +28 -25
  75. ultralytics/models/yolo/pose/val.py +91 -64
  76. ultralytics/models/yolo/segment/__init__.py +1 -1
  77. ultralytics/models/yolo/segment/predict.py +10 -8
  78. ultralytics/models/yolo/segment/train.py +16 -15
  79. ultralytics/models/yolo/segment/val.py +90 -68
  80. ultralytics/nn/__init__.py +26 -6
  81. ultralytics/nn/autobackend.py +144 -112
  82. ultralytics/nn/modules/__init__.py +96 -13
  83. ultralytics/nn/modules/block.py +28 -7
  84. ultralytics/nn/modules/conv.py +41 -23
  85. ultralytics/nn/modules/head.py +67 -59
  86. ultralytics/nn/modules/transformer.py +49 -32
  87. ultralytics/nn/modules/utils.py +20 -15
  88. ultralytics/nn/tasks.py +215 -141
  89. ultralytics/solutions/ai_gym.py +59 -47
  90. ultralytics/solutions/distance_calculation.py +22 -15
  91. ultralytics/solutions/heatmap.py +76 -54
  92. ultralytics/solutions/object_counter.py +46 -39
  93. ultralytics/solutions/speed_estimation.py +13 -16
  94. ultralytics/trackers/__init__.py +1 -1
  95. ultralytics/trackers/basetrack.py +1 -0
  96. ultralytics/trackers/bot_sort.py +2 -1
  97. ultralytics/trackers/byte_tracker.py +10 -7
  98. ultralytics/trackers/track.py +7 -7
  99. ultralytics/trackers/utils/gmc.py +25 -25
  100. ultralytics/trackers/utils/kalman_filter.py +85 -42
  101. ultralytics/trackers/utils/matching.py +8 -7
  102. ultralytics/utils/__init__.py +173 -151
  103. ultralytics/utils/autobatch.py +10 -10
  104. ultralytics/utils/benchmarks.py +76 -86
  105. ultralytics/utils/callbacks/__init__.py +1 -1
  106. ultralytics/utils/callbacks/base.py +29 -29
  107. ultralytics/utils/callbacks/clearml.py +51 -43
  108. ultralytics/utils/callbacks/comet.py +81 -66
  109. ultralytics/utils/callbacks/dvc.py +33 -26
  110. ultralytics/utils/callbacks/hub.py +44 -26
  111. ultralytics/utils/callbacks/mlflow.py +31 -24
  112. ultralytics/utils/callbacks/neptune.py +35 -25
  113. ultralytics/utils/callbacks/raytune.py +9 -4
  114. ultralytics/utils/callbacks/tensorboard.py +16 -11
  115. ultralytics/utils/callbacks/wb.py +39 -33
  116. ultralytics/utils/checks.py +189 -141
  117. ultralytics/utils/dist.py +15 -12
  118. ultralytics/utils/downloads.py +112 -96
  119. ultralytics/utils/errors.py +1 -1
  120. ultralytics/utils/files.py +11 -11
  121. ultralytics/utils/instance.py +22 -22
  122. ultralytics/utils/loss.py +117 -67
  123. ultralytics/utils/metrics.py +224 -158
  124. ultralytics/utils/ops.py +39 -29
  125. ultralytics/utils/patches.py +3 -3
  126. ultralytics/utils/plotting.py +217 -120
  127. ultralytics/utils/tal.py +19 -13
  128. ultralytics/utils/torch_utils.py +138 -109
  129. ultralytics/utils/triton.py +12 -10
  130. ultralytics/utils/tuner.py +49 -47
  131. {ultralytics-8.0.237.dist-info → ultralytics-8.0.239.dist-info}/METADATA +5 -4
  132. ultralytics-8.0.239.dist-info/RECORD +188 -0
  133. ultralytics-8.0.237.dist-info/RECORD +0 -187
  134. {ultralytics-8.0.237.dist-info → ultralytics-8.0.239.dist-info}/LICENSE +0 -0
  135. {ultralytics-8.0.237.dist-info → ultralytics-8.0.239.dist-info}/WHEEL +0 -0
  136. {ultralytics-8.0.237.dist-info → ultralytics-8.0.239.dist-info}/entry_points.txt +0 -0
  137. {ultralytics-8.0.237.dist-info → ultralytics-8.0.239.dist-info}/top_level.txt +0 -0
@@ -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 = 'pushup'
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(self,
40
- kpts_to_check,
41
- line_thickness=2,
42
- view_img=False,
43
- pose_up_angle=145.0,
44
- pose_down_angle=90.0,
45
- pose_type='pullup'):
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 = ['-' for _ in results[0]]
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 == 'pushup' or self.pose_type == 'pullup':
81
- self.angle[ind] = self.annotator.estimate_pose_angle(k[int(self.kpts_to_check[0])].cpu(),
82
- k[int(self.kpts_to_check[1])].cpu(),
83
- k[int(self.kpts_to_check[2])].cpu())
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 == 'abworkout':
87
- self.angle[ind] = self.annotator.estimate_pose_angle(k[int(self.kpts_to_check[0])].cpu(),
88
- k[int(self.kpts_to_check[1])].cpu(),
89
- k[int(self.kpts_to_check[2])].cpu())
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] = 'down'
93
- if self.angle[ind] < self.posedown_angle and self.stage[ind] == 'down':
94
- self.stage[ind] = 'up'
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(angle_text=self.angle[ind],
97
- count_text=self.count[ind],
98
- stage_text=self.stage[ind],
99
- center_kpt=k[int(self.kpts_to_check[1])],
100
- line_thickness=self.tf)
101
-
102
- if self.pose_type == 'pushup':
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] = 'up'
105
- if self.angle[ind] < self.posedown_angle and self.stage[ind] == 'up':
106
- self.stage[ind] = 'down'
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(angle_text=self.angle[ind],
109
- count_text=self.count[ind],
110
- stage_text=self.stage[ind],
111
- center_kpt=k[int(self.kpts_to_check[1])],
112
- line_thickness=self.tf)
113
- if self.pose_type == 'pullup':
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] = 'down'
116
- if self.angle[ind] < self.posedown_angle and self.stage[ind] == 'down':
117
- self.stage[ind] = 'up'
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(angle_text=self.angle[ind],
120
- count_text=self.count[ind],
121
- stage_text=self.stage[ind],
122
- center_kpt=k[int(self.kpts_to_check[1])],
123
- line_thickness=self.tf)
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('Ultralytics YOLOv8 AI GYM', self.im0)
129
- if cv2.waitKey(1) & 0xFF == ord('q'):
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__ == '__main__':
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
- def set_args(self,
41
- names,
42
- pixels_per_meter=10,
43
- view_img=False,
44
- line_thickness=2,
45
- line_color=(255, 255, 0),
46
- centroid_color=(255, 0, 255)):
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(self.im0, f'Distance : {distance:.2f}m', (20, 55), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 0), 2,
129
- cv2.LINE_AA)
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('Ultralytics Distance Estimation')
179
- cv2.setMouseCallback('Ultralytics Distance Estimation', self.mouse_event_for_distance)
180
- cv2.imshow('Ultralytics Distance Estimation', self.im0)
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('q'):
189
+ if cv2.waitKey(1) & 0xFF == ord("q"):
183
190
  return
184
191
 
185
192
 
186
- if __name__ == '__main__':
193
+ if __name__ == "__main__":
187
194
  DistanceCalculation()
@@ -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('shapely>=2.0.0')
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 = 'circle'
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(self,
65
- imw,
66
- imh,
67
- colormap=cv2.COLORMAP_JET,
68
- heatmap_alpha=0.5,
69
- view_img=False,
70
- count_reg_pts=None,
71
- count_txt_thickness=2,
72
- count_txt_color=(0, 0, 0),
73
- count_color=(255, 255, 255),
74
- count_reg_color=(255, 0, 255),
75
- region_thickness=5,
76
- line_dist_thresh=15,
77
- decay_factor=0.99,
78
- shape='circle'):
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('Line Counter Initiated.')
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('Region Counter Initiated.')
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('Region or line points Invalid, 2 or 4 points supported')
119
- print('Using Line Counter Now')
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 ['circle', 'rect']:
144
+ if self.shape not in ["circle", "rect"]:
136
145
  print("Unknown shape value provided, 'circle' & 'rect' supported")
137
- print('Using Circular shape now')
138
- self.shape = 'circle'
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.annotator.draw_region(reg_pts=self.count_reg_pts,
175
- color=self.region_color,
176
- thickness=self.region_thickness)
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 ** 2
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
- (2 * mask[int(box[1]):int(box[3]), int(box[0]):int(box[2])])
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 ** 2
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
- (2 * mask[int(box[1]):int(box[3]), int(box[0]):int(box[2])])
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
- if self.count_reg_pts is not None:
239
- incount_label = 'InCount : ' + f'{self.in_counts}'
240
- outcount_label = 'OutCount : ' + f'{self.out_counts}'
241
- self.annotator.count_labels(in_count=incount_label,
242
- out_count=outcount_label,
243
- count_txt_size=self.count_txt_thickness,
244
- txt_color=self.count_txt_color,
245
- color=self.count_color)
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('Ultralytics Heatmap', self.im0)
278
+ cv2.imshow("Ultralytics Heatmap", self.im0)
257
279
 
258
- if cv2.waitKey(1) & 0xFF == ord('q'):
280
+ if cv2.waitKey(1) & 0xFF == ord("q"):
259
281
  return
260
282
 
261
283
 
262
- if __name__ == '__main__':
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('shapely>=2.0.0')
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(self,
60
- classes_names,
61
- reg_pts,
62
- count_reg_color=(255, 0, 255),
63
- line_thickness=2,
64
- track_thickness=2,
65
- view_img=False,
66
- view_in_counts=True,
67
- view_out_counts=True,
68
- draw_tracks=False,
69
- count_txt_thickness=2,
70
- count_txt_color=(0, 0, 0),
71
- count_color=(255, 255, 255),
72
- track_color=(0, 255, 0),
73
- region_thickness=5,
74
- line_dist_thresh=15):
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('Line Counter Initiated.')
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('Region Counter Initiated.')
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('Invalid Region points provided, region_points can be 2 or 4')
113
- print('Using Line Counter Now')
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(box, label=str(track_id) + ':' + self.names[cls],
167
- color=colors(int(cls), True)) # Draw bounding box
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(track_line,
178
- color=self.track_color,
179
- track_thickness=self.track_thickness)
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 = 'In Count : ' + f'{self.in_counts}'
202
- outcount_label = 'OutCount : ' + f'{self.out_counts}'
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 + ' ' + outcount_label
217
+ counts_label = incount_label + " " + outcount_label
214
218
 
215
219
  if counts_label is not None:
216
- self.annotator.count_labels(counts=counts_label,
217
- count_txt_size=self.count_txt_thickness,
218
- txt_color=self.count_txt_color,
219
- color=self.count_color)
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('Ultralytics YOLOv8 Object Counter')
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('Ultralytics YOLOv8 Object Counter', self.mouse_event_for_region,
227
- {'region_points': self.reg_pts})
228
- cv2.imshow('Ultralytics YOLOv8 Object Counter', self.im0)
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('q'):
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__ == '__main__':
263
+ if __name__ == "__main__":
257
264
  ObjectCounter()