ultralytics 8.1.37__py3-none-any.whl → 8.1.39__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 (39) hide show
  1. ultralytics/__init__.py +1 -1
  2. ultralytics/cfg/__init__.py +1 -2
  3. ultralytics/cfg/datasets/lvis.yaml +1239 -0
  4. ultralytics/cfg/default.yaml +2 -2
  5. ultralytics/data/__init__.py +18 -2
  6. ultralytics/data/augment.py +123 -2
  7. ultralytics/data/base.py +2 -0
  8. ultralytics/data/build.py +25 -3
  9. ultralytics/data/converter.py +22 -4
  10. ultralytics/data/dataset.py +143 -27
  11. ultralytics/data/utils.py +25 -1
  12. ultralytics/engine/exporter.py +1 -3
  13. ultralytics/engine/model.py +4 -1
  14. ultralytics/engine/trainer.py +48 -44
  15. ultralytics/models/fastsam/prompt.py +1 -1
  16. ultralytics/models/yolo/__init__.py +2 -2
  17. ultralytics/models/yolo/detect/val.py +36 -17
  18. ultralytics/models/yolo/model.py +1 -0
  19. ultralytics/models/yolo/world/__init__.py +5 -0
  20. ultralytics/models/yolo/world/train.py +91 -0
  21. ultralytics/models/yolo/world/train_world.py +108 -0
  22. ultralytics/nn/autobackend.py +1 -1
  23. ultralytics/nn/modules/block.py +4 -2
  24. ultralytics/nn/modules/head.py +9 -0
  25. ultralytics/nn/tasks.py +29 -13
  26. ultralytics/solutions/heatmap.py +84 -46
  27. ultralytics/solutions/object_counter.py +79 -64
  28. ultralytics/trackers/utils/gmc.py +1 -1
  29. ultralytics/utils/callbacks/raytune.py +1 -1
  30. ultralytics/utils/loss.py +1 -1
  31. ultralytics/utils/plotting.py +35 -21
  32. ultralytics/utils/torch_utils.py +14 -0
  33. ultralytics/utils/tuner.py +2 -2
  34. {ultralytics-8.1.37.dist-info → ultralytics-8.1.39.dist-info}/METADATA +1 -1
  35. {ultralytics-8.1.37.dist-info → ultralytics-8.1.39.dist-info}/RECORD +39 -35
  36. {ultralytics-8.1.37.dist-info → ultralytics-8.1.39.dist-info}/LICENSE +0 -0
  37. {ultralytics-8.1.37.dist-info → ultralytics-8.1.39.dist-info}/WHEEL +0 -0
  38. {ultralytics-8.1.37.dist-info → ultralytics-8.1.39.dist-info}/entry_points.txt +0 -0
  39. {ultralytics-8.1.37.dist-info → ultralytics-8.1.39.dist-info}/top_level.txt +0 -0
@@ -24,6 +24,8 @@ class Heatmap:
24
24
  self.view_img = False
25
25
  self.shape = "circle"
26
26
 
27
+ self.names = None # Classes names
28
+
27
29
  # Image information
28
30
  self.imw = None
29
31
  self.imh = None
@@ -52,10 +54,13 @@ class Heatmap:
52
54
  # Object Counting Information
53
55
  self.in_counts = 0
54
56
  self.out_counts = 0
55
- self.counting_list = []
57
+ self.count_ids = []
58
+ self.class_wise_count = {}
56
59
  self.count_txt_thickness = 0
57
- self.count_txt_color = (0, 0, 0)
58
- self.count_color = (255, 255, 255)
60
+ self.count_txt_color = (255, 255, 255)
61
+ self.line_color = (255, 255, 255)
62
+ self.cls_txtdisplay_gap = 50
63
+ self.fontsize = 0.6
59
64
 
60
65
  # Decay factor
61
66
  self.decay_factor = 0.99
@@ -67,6 +72,7 @@ class Heatmap:
67
72
  self,
68
73
  imw,
69
74
  imh,
75
+ classes_names=None,
70
76
  colormap=cv2.COLORMAP_JET,
71
77
  heatmap_alpha=0.5,
72
78
  view_img=False,
@@ -74,13 +80,15 @@ class Heatmap:
74
80
  view_out_counts=True,
75
81
  count_reg_pts=None,
76
82
  count_txt_thickness=2,
77
- count_txt_color=(0, 0, 0),
78
- count_color=(255, 255, 255),
83
+ count_txt_color=(255, 255, 255),
84
+ fontsize=0.8,
85
+ line_color=(255, 255, 255),
79
86
  count_reg_color=(255, 0, 255),
80
87
  region_thickness=5,
81
88
  line_dist_thresh=15,
82
89
  decay_factor=0.99,
83
90
  shape="circle",
91
+ cls_txtdisplay_gap=50,
84
92
  ):
85
93
  """
86
94
  Configures the heatmap colormap, width, height and display parameters.
@@ -89,6 +97,7 @@ class Heatmap:
89
97
  colormap (cv2.COLORMAP): The colormap to be set.
90
98
  imw (int): The width of the frame.
91
99
  imh (int): The height of the frame.
100
+ classes_names (dict): Classes names
92
101
  heatmap_alpha (float): alpha value for heatmap display
93
102
  view_img (bool): Flag indicating frame display
94
103
  view_in_counts (bool): Flag to control whether to display the incounts on video stream.
@@ -96,13 +105,16 @@ class Heatmap:
96
105
  count_reg_pts (list): Object counting region points
97
106
  count_txt_thickness (int): Text thickness for object counting display
98
107
  count_txt_color (RGB color): count text color value
99
- count_color (RGB color): count text background color value
108
+ fontsize (float): Text display font size
109
+ line_color (RGB color): count highlighter line color
100
110
  count_reg_color (RGB color): Color of object counting region
101
111
  region_thickness (int): Object counting Region thickness
102
112
  line_dist_thresh (int): Euclidean Distance threshold for line counter
103
113
  decay_factor (float): value for removing heatmap area after object passed
104
114
  shape (str): Heatmap shape, rect or circle shape supported
115
+ cls_txtdisplay_gap (int): Display gap between each class count
105
116
  """
117
+ self.names = classes_names
106
118
  self.imw = imw
107
119
  self.imh = imh
108
120
  self.heatmap_alpha = heatmap_alpha
@@ -116,29 +128,29 @@ class Heatmap:
116
128
  if len(count_reg_pts) == 2:
117
129
  print("Line Counter Initiated.")
118
130
  self.count_reg_pts = count_reg_pts
119
- self.counting_region = LineString(count_reg_pts)
120
-
121
- elif len(count_reg_pts) == 4:
122
- print("Region Counter Initiated.")
131
+ self.counting_region = LineString(self.count_reg_pts)
132
+ elif len(count_reg_pts) >= 3:
133
+ print("Polygon Counter Initiated.")
123
134
  self.count_reg_pts = count_reg_pts
124
135
  self.counting_region = Polygon(self.count_reg_pts)
125
-
126
136
  else:
127
- print("Region or line points Invalid, 2 or 4 points supported")
137
+ print("Invalid Region points provided, region_points must be 2 for lines or >= 3 for polygons.")
128
138
  print("Using Line Counter Now")
129
- self.counting_region = Polygon([(20, 400), (1260, 400)]) # dummy points
139
+ self.counting_region = LineString(self.count_reg_pts)
130
140
 
131
141
  # Heatmap new frame
132
142
  self.heatmap = np.zeros((int(self.imh), int(self.imw)), dtype=np.float32)
133
143
 
134
144
  self.count_txt_thickness = count_txt_thickness
135
145
  self.count_txt_color = count_txt_color
136
- self.count_color = count_color
146
+ self.fontsize = fontsize
147
+ self.line_color = line_color
137
148
  self.region_color = count_reg_color
138
149
  self.region_thickness = region_thickness
139
150
  self.decay_factor = decay_factor
140
151
  self.line_dist_thresh = line_dist_thresh
141
152
  self.shape = shape
153
+ self.cls_txtdisplay_gap = cls_txtdisplay_gap
142
154
 
143
155
  # shape of heatmap, if not selected
144
156
  if self.shape not in ["circle", "rect"]:
@@ -183,6 +195,12 @@ class Heatmap:
183
195
  )
184
196
 
185
197
  for box, cls, track_id in zip(self.boxes, self.clss, self.track_ids):
198
+ # Store class info
199
+ if self.names[cls] not in self.class_wise_count:
200
+ if len(self.names[cls]) > 5:
201
+ self.names[cls] = self.names[cls][:5]
202
+ self.class_wise_count[self.names[cls]] = {"in": 0, "out": 0}
203
+
186
204
  if self.shape == "circle":
187
205
  center = (int((box[0] + box[2]) // 2), int((box[1] + box[3]) // 2))
188
206
  radius = min(int(box[2]) - int(box[0]), int(box[3]) - int(box[1])) // 2
@@ -203,23 +221,39 @@ class Heatmap:
203
221
  if len(track_line) > 30:
204
222
  track_line.pop(0)
205
223
 
206
- # Count objects
207
- if len(self.count_reg_pts) == 4:
208
- if self.counting_region.contains(Point(track_line[-1])) and track_id not in self.counting_list:
209
- self.counting_list.append(track_id)
210
- if box[0] < self.counting_region.centroid.x:
211
- self.out_counts += 1
212
- else:
224
+ prev_position = self.track_history[track_id][-2] if len(self.track_history[track_id]) > 1 else None
225
+
226
+ # Count objects in any polygon
227
+ if len(self.count_reg_pts) >= 3:
228
+ is_inside = self.counting_region.contains(Point(track_line[-1]))
229
+
230
+ if prev_position is not None and is_inside and track_id not in self.count_ids:
231
+ self.count_ids.append(track_id)
232
+
233
+ if (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0:
213
234
  self.in_counts += 1
235
+ self.class_wise_count[self.names[cls]]["in"] += 1
236
+ else:
237
+ self.out_counts += 1
238
+ self.class_wise_count[self.names[cls]]["out"] += 1
214
239
 
240
+ # Count objects using line
215
241
  elif len(self.count_reg_pts) == 2:
216
- distance = Point(track_line[-1]).distance(self.counting_region)
217
- if distance < self.line_dist_thresh and track_id not in self.counting_list:
218
- self.counting_list.append(track_id)
219
- if box[0] < self.counting_region.centroid.x:
220
- self.out_counts += 1
221
- else:
222
- self.in_counts += 1
242
+ is_inside = (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0
243
+
244
+ if prev_position is not None and is_inside and track_id not in self.count_ids:
245
+ distance = Point(track_line[-1]).distance(self.counting_region)
246
+
247
+ if distance < self.line_dist_thresh and track_id not in self.count_ids:
248
+ self.count_ids.append(track_id)
249
+
250
+ if (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0:
251
+ self.in_counts += 1
252
+ self.class_wise_count[self.names[cls]]["in"] += 1
253
+ else:
254
+ self.out_counts += 1
255
+ self.class_wise_count[self.names[cls]]["out"] += 1
256
+
223
257
  else:
224
258
  for box, cls in zip(self.boxes, self.clss):
225
259
  if self.shape == "circle":
@@ -240,26 +274,30 @@ class Heatmap:
240
274
  heatmap_normalized = cv2.normalize(self.heatmap, None, 0, 255, cv2.NORM_MINMAX)
241
275
  heatmap_colored = cv2.applyColorMap(heatmap_normalized.astype(np.uint8), self.colormap)
242
276
 
243
- incount_label = f"In Count : {self.in_counts}"
244
- outcount_label = f"OutCount : {self.out_counts}"
245
-
246
- # Display counts based on user choice
247
- counts_label = None
248
- if not self.view_in_counts and not self.view_out_counts:
249
- counts_label = None
250
- elif not self.view_in_counts:
251
- counts_label = outcount_label
252
- elif not self.view_out_counts:
253
- counts_label = incount_label
254
- else:
255
- counts_label = f"{incount_label} {outcount_label}"
277
+ label = "Ultralytics Analytics \t"
278
+
279
+ for key, value in self.class_wise_count.items():
280
+ if value["in"] != 0 or value["out"] != 0:
281
+ if not self.view_in_counts and not self.view_out_counts:
282
+ label = None
283
+ elif not self.view_in_counts:
284
+ label += f"{str.capitalize(key)}: IN {value['in']} \t"
285
+ elif not self.view_out_counts:
286
+ label += f"{str.capitalize(key)}: OUT {value['out']} \t"
287
+ else:
288
+ label += f"{str.capitalize(key)}: IN {value['in']} OUT {value['out']} \t"
289
+
290
+ label = label.rstrip()
291
+ label = label.split("\t")
256
292
 
257
- if self.count_reg_pts is not None and counts_label is not None:
258
- self.annotator.count_labels(
259
- counts=counts_label,
260
- count_txt_size=self.count_txt_thickness,
293
+ if self.count_reg_pts is not None and label is not None:
294
+ self.annotator.display_counts(
295
+ counts=label,
296
+ tf=self.count_txt_thickness,
297
+ fontScale=self.fontsize,
261
298
  txt_color=self.count_txt_color,
262
- color=self.count_color,
299
+ line_color=self.line_color,
300
+ classwise_txtgap=self.cls_txtdisplay_gap,
263
301
  )
264
302
 
265
303
  self.im0 = cv2.addWeighted(self.im0, 1 - self.heatmap_alpha, heatmap_colored, self.heatmap_alpha, 0)
@@ -43,16 +43,19 @@ class ObjectCounter:
43
43
  # Object counting Information
44
44
  self.in_counts = 0
45
45
  self.out_counts = 0
46
- self.counting_dict = {}
46
+ self.count_ids = []
47
+ self.class_wise_count = {}
47
48
  self.count_txt_thickness = 0
48
- self.count_txt_color = (0, 0, 0)
49
- self.count_color = (255, 255, 255)
49
+ self.count_txt_color = (255, 255, 255)
50
+ self.line_color = (255, 255, 255)
51
+ self.cls_txtdisplay_gap = 50
52
+ self.fontsize = 0.6
50
53
 
51
54
  # Tracks info
52
55
  self.track_history = defaultdict(list)
53
56
  self.track_thickness = 2
54
57
  self.draw_tracks = False
55
- self.track_color = (0, 255, 0)
58
+ self.track_color = None
56
59
 
57
60
  # Check if environment support imshow
58
61
  self.env_check = check_imshow(warn=True)
@@ -68,12 +71,14 @@ class ObjectCounter:
68
71
  view_in_counts=True,
69
72
  view_out_counts=True,
70
73
  draw_tracks=False,
71
- count_txt_thickness=2,
72
- count_txt_color=(0, 0, 0),
73
- count_color=(255, 255, 255),
74
- track_color=(0, 255, 0),
74
+ count_txt_thickness=3,
75
+ count_txt_color=(255, 255, 255),
76
+ fontsize=0.8,
77
+ line_color=(255, 255, 255),
78
+ track_color=None,
75
79
  region_thickness=5,
76
80
  line_dist_thresh=15,
81
+ cls_txtdisplay_gap=50,
77
82
  ):
78
83
  """
79
84
  Configures the Counter's image, bounding box line thickness, and counting region points.
@@ -89,11 +94,13 @@ class ObjectCounter:
89
94
  draw_tracks (Bool): draw tracks
90
95
  count_txt_thickness (int): Text thickness for object counting display
91
96
  count_txt_color (RGB color): count text color value
92
- count_color (RGB color): count text background color value
97
+ fontsize (float): Text display font size
98
+ line_color (RGB color): count highlighter line color
93
99
  count_reg_color (RGB color): Color of object counting region
94
100
  track_color (RGB color): color for tracks
95
101
  region_thickness (int): Object counting Region thickness
96
102
  line_dist_thresh (int): Euclidean Distance threshold for line counter
103
+ cls_txtdisplay_gap (int): Display gap between each class count
97
104
  """
98
105
  self.tf = line_thickness
99
106
  self.view_img = view_img
@@ -108,7 +115,7 @@ class ObjectCounter:
108
115
  self.reg_pts = reg_pts
109
116
  self.counting_region = LineString(self.reg_pts)
110
117
  elif len(reg_pts) >= 3:
111
- print("Region Counter Initiated.")
118
+ print("Polygon Counter Initiated.")
112
119
  self.reg_pts = reg_pts
113
120
  self.counting_region = Polygon(self.reg_pts)
114
121
  else:
@@ -120,10 +127,12 @@ class ObjectCounter:
120
127
  self.track_color = track_color
121
128
  self.count_txt_thickness = count_txt_thickness
122
129
  self.count_txt_color = count_txt_color
123
- self.count_color = count_color
130
+ self.fontsize = fontsize
131
+ self.line_color = line_color
124
132
  self.region_color = count_reg_color
125
133
  self.region_thickness = region_thickness
126
134
  self.line_dist_thresh = line_dist_thresh
135
+ self.cls_txtdisplay_gap = cls_txtdisplay_gap
127
136
 
128
137
  def mouse_event_for_region(self, event, x, y, flags, params):
129
138
  """
@@ -171,7 +180,13 @@ class ObjectCounter:
171
180
  # Extract tracks
172
181
  for box, track_id, cls in zip(boxes, track_ids, clss):
173
182
  # Draw bounding box
174
- self.annotator.box_label(box, label=f"{track_id}:{self.names[cls]}", color=colors(int(track_id), True))
183
+ self.annotator.box_label(box, label=f"{self.names[cls]}#{track_id}", color=colors(int(track_id), True))
184
+
185
+ # Store class info
186
+ if self.names[cls] not in self.class_wise_count:
187
+ if len(self.names[cls]) > 5:
188
+ self.names[cls] = self.names[cls][:5]
189
+ self.class_wise_count[self.names[cls]] = {"in": 0, "out": 0}
175
190
 
176
191
  # Draw Tracks
177
192
  track_line = self.track_history[track_id]
@@ -182,68 +197,68 @@ class ObjectCounter:
182
197
  # Draw track trails
183
198
  if self.draw_tracks:
184
199
  self.annotator.draw_centroid_and_tracks(
185
- track_line, color=self.track_color, track_thickness=self.track_thickness
200
+ track_line,
201
+ color=self.track_color if self.track_color else colors(int(track_id), True),
202
+ track_thickness=self.track_thickness,
186
203
  )
187
204
 
188
205
  prev_position = self.track_history[track_id][-2] if len(self.track_history[track_id]) > 1 else None
189
- centroid = Point((box[:2] + box[2:]) / 2)
190
-
191
- # Count objects
192
- if len(self.reg_pts) >= 3: # any polygon
193
- is_inside = self.counting_region.contains(centroid)
194
- current_position = "in" if is_inside else "out"
195
206
 
196
- if prev_position is not None:
197
- if self.counting_dict[track_id] != current_position and is_inside:
198
- self.in_counts += 1
199
- self.counting_dict[track_id] = "in"
200
- elif self.counting_dict[track_id] != current_position and not is_inside:
201
- self.out_counts += 1
202
- self.counting_dict[track_id] = "out"
203
- else:
204
- self.counting_dict[track_id] = current_position
207
+ # Count objects in any polygon
208
+ if len(self.reg_pts) >= 3:
209
+ is_inside = self.counting_region.contains(Point(track_line[-1]))
205
210
 
206
- else:
207
- self.counting_dict[track_id] = current_position
211
+ if prev_position is not None and is_inside and track_id not in self.count_ids:
212
+ self.count_ids.append(track_id)
208
213
 
209
- elif len(self.reg_pts) == 2:
210
- if prev_position is not None:
211
- is_inside = (box[0] - prev_position[0]) * (
212
- self.counting_region.centroid.x - prev_position[0]
213
- ) > 0
214
- current_position = "in" if is_inside else "out"
215
-
216
- if self.counting_dict[track_id] != current_position and is_inside:
214
+ if (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0:
217
215
  self.in_counts += 1
218
- self.counting_dict[track_id] = "in"
219
- elif self.counting_dict[track_id] != current_position and not is_inside:
220
- self.out_counts += 1
221
- self.counting_dict[track_id] = "out"
216
+ self.class_wise_count[self.names[cls]]["in"] += 1
222
217
  else:
223
- self.counting_dict[track_id] = current_position
224
- else:
225
- self.counting_dict[track_id] = None
226
-
227
- incount_label = f"In Count : {self.in_counts}"
228
- outcount_label = f"OutCount : {self.out_counts}"
229
-
230
- # Display counts based on user choice
231
- counts_label = None
232
- if not self.view_in_counts and not self.view_out_counts:
233
- counts_label = None
234
- elif not self.view_in_counts:
235
- counts_label = outcount_label
236
- elif not self.view_out_counts:
237
- counts_label = incount_label
238
- else:
239
- counts_label = f"{incount_label} {outcount_label}"
218
+ self.out_counts += 1
219
+ self.class_wise_count[self.names[cls]]["out"] += 1
240
220
 
241
- if counts_label is not None:
242
- self.annotator.count_labels(
243
- counts=counts_label,
244
- count_txt_size=self.count_txt_thickness,
221
+ # Count objects using line
222
+ elif len(self.reg_pts) == 2:
223
+ is_inside = (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0
224
+
225
+ if prev_position is not None and is_inside and track_id not in self.count_ids:
226
+ distance = Point(track_line[-1]).distance(self.counting_region)
227
+
228
+ if distance < self.line_dist_thresh and track_id not in self.count_ids:
229
+ self.count_ids.append(track_id)
230
+
231
+ if (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0:
232
+ self.in_counts += 1
233
+ self.class_wise_count[self.names[cls]]["in"] += 1
234
+ else:
235
+ self.out_counts += 1
236
+ self.class_wise_count[self.names[cls]]["out"] += 1
237
+
238
+ label = "Ultralytics Analytics \t"
239
+
240
+ for key, value in self.class_wise_count.items():
241
+ if value["in"] != 0 or value["out"] != 0:
242
+ if not self.view_in_counts and not self.view_out_counts:
243
+ label = None
244
+ elif not self.view_in_counts:
245
+ label += f"{str.capitalize(key)}: IN {value['in']} \t"
246
+ elif not self.view_out_counts:
247
+ label += f"{str.capitalize(key)}: OUT {value['out']} \t"
248
+ else:
249
+ label += f"{str.capitalize(key)}: IN {value['in']} OUT {value['out']} \t"
250
+
251
+ label = label.rstrip()
252
+ label = label.split("\t")
253
+
254
+ if label is not None:
255
+ self.annotator.display_counts(
256
+ counts=label,
257
+ tf=self.count_txt_thickness,
258
+ fontScale=self.fontsize,
245
259
  txt_color=self.count_txt_color,
246
- color=self.count_color,
260
+ line_color=self.line_color,
261
+ classwise_txtgap=self.cls_txtdisplay_gap,
247
262
  )
248
263
 
249
264
  def display_frames(self):
@@ -319,7 +319,7 @@ class GMC:
319
319
  keypoints = cv2.goodFeaturesToTrack(frame, mask=None, **self.feature_params)
320
320
 
321
321
  # Handle first frame
322
- if not self.initializedFirstFrame:
322
+ if not self.initializedFirstFrame or self.prevKeyPoints is None:
323
323
  self.prevFrame = frame.copy()
324
324
  self.prevKeyPoints = copy.copy(keypoints)
325
325
  self.initializedFirstFrame = True
@@ -14,7 +14,7 @@ except (ImportError, AssertionError):
14
14
 
15
15
  def on_fit_epoch_end(trainer):
16
16
  """Sends training metrics to Ray Tune at end of each epoch."""
17
- if ray.tune.is_session_enabled():
17
+ if ray.train._internal.session._get_session(): # replacement for deprecated ray.tune.is_session_enabled()
18
18
  metrics = trainer.metrics
19
19
  metrics["epoch"] = trainer.epoch
20
20
  session.report(metrics)
ultralytics/utils/loss.py CHANGED
@@ -157,7 +157,7 @@ class v8DetectionLoss:
157
157
  self.hyp = h
158
158
  self.stride = m.stride # model strides
159
159
  self.nc = m.nc # number of classes
160
- self.no = m.no
160
+ self.no = m.nc + m.reg_max * 4
161
161
  self.reg_max = m.reg_max
162
162
  self.device = device
163
163
 
@@ -363,35 +363,49 @@ class Annotator:
363
363
  cv2.polylines(self.im, [points], isClosed=False, color=color, thickness=track_thickness)
364
364
  cv2.circle(self.im, (int(track[-1][0]), int(track[-1][1])), track_thickness * 2, color, -1)
365
365
 
366
- def count_labels(self, counts=0, count_txt_size=2, color=(255, 255, 255), txt_color=(0, 0, 0)):
366
+ def display_counts(
367
+ self, counts=None, tf=2, fontScale=0.6, line_color=(0, 0, 0), txt_color=(255, 255, 255), classwise_txtgap=55
368
+ ):
367
369
  """
368
- Plot counts for object counter.
370
+ Display counts on im0.
369
371
 
370
372
  Args:
371
- counts (int): objects counts value
372
- count_txt_size (int): text size for counts display
373
- color (tuple): background color of counts display
374
- txt_color (tuple): text color of counts display
373
+ counts (str): objects count data
374
+ tf (int): text thickness for display
375
+ fontScale (float): text fontsize for display
376
+ line_color (RGB Color): counts highlighter color
377
+ txt_color (RGB Color): counts display color
378
+ classwise_txtgap (int): Gap between each class count data
375
379
  """
376
- self.tf = count_txt_size
377
- tl = self.tf or round(0.002 * (self.im.shape[0] + self.im.shape[1]) / 2) + 1
380
+
381
+ tl = tf or round(0.002 * (self.im.shape[0] + self.im.shape[1]) / 2) + 1
378
382
  tf = max(tl - 1, 1)
379
383
 
380
- # Get text size for in_count and out_count
381
- t_size_in = cv2.getTextSize(str(counts), 0, fontScale=tl / 2, thickness=tf)[0]
384
+ t_sizes = [cv2.getTextSize(str(count), 0, fontScale=0.8, thickness=tf)[0] for count in counts]
382
385
 
383
- # Calculate positions for counts label
384
- text_width = t_size_in[0]
385
- text_x = (self.im.shape[1] - text_width) // 2 # Center x-coordinate
386
- text_y = t_size_in[1]
386
+ max_text_width = max([size[0] for size in t_sizes])
387
+ max_text_height = max([size[1] for size in t_sizes])
387
388
 
388
- # Create a rounded rectangle for in_count
389
- cv2.rectangle(
390
- self.im, (text_x - 5, text_y - 5), (text_x + text_width + 7, text_y + t_size_in[1] + 7), color, -1
391
- )
392
- cv2.putText(
393
- self.im, str(counts), (text_x, text_y + t_size_in[1]), 0, tl / 2, txt_color, self.tf, lineType=cv2.LINE_AA
394
- )
389
+ text_x = self.im.shape[1] - max_text_width - 20
390
+ text_y = classwise_txtgap
391
+
392
+ for i, count in enumerate(counts):
393
+ text_x_pos = text_x
394
+ text_y_pos = text_y + i * classwise_txtgap
395
+
396
+ cv2.putText(
397
+ self.im,
398
+ str(count),
399
+ (text_x_pos, text_y_pos),
400
+ cv2.FONT_HERSHEY_SIMPLEX,
401
+ fontScale=fontScale,
402
+ color=txt_color,
403
+ thickness=tf,
404
+ lineType=cv2.LINE_AA,
405
+ )
406
+
407
+ line_y_pos = text_y_pos + max_text_height + 5
408
+ cv2.line(self.im, (text_x_pos, line_y_pos), (text_x_pos + max_text_width, line_y_pos), line_color, tf)
395
409
 
396
410
  @staticmethod
397
411
  def estimate_pose_angle(a, b, c):
@@ -505,6 +505,20 @@ def strip_optimizer(f: Union[str, Path] = "best.pt", s: str = "") -> None:
505
505
  LOGGER.info(f"Optimizer stripped from {f},{f' saved as {s},' if s else ''} {mb:.1f}MB")
506
506
 
507
507
 
508
+ def convert_optimizer_state_dict_to_fp16(state_dict):
509
+ """
510
+ Converts the state_dict of a given optimizer to FP16, focusing on the 'state' key for tensor conversions.
511
+
512
+ This method aims to reduce storage size without altering 'param_groups' as they contain non-tensor data.
513
+ """
514
+ for state in state_dict["state"].values():
515
+ for k, v in state.items():
516
+ if isinstance(v, torch.Tensor) and v.dtype is torch.float32:
517
+ state[k] = v.half()
518
+
519
+ return state_dict
520
+
521
+
508
522
  def profile(input, ops, n=10, device=None):
509
523
  """
510
524
  Ultralytics speed, memory and FLOPs profiler.
@@ -40,7 +40,7 @@ def run_ray_tune(
40
40
  train_args = {}
41
41
 
42
42
  try:
43
- subprocess.run("pip install ray[tune]<=2.9.3".split(), check=True) # do not add single quotes here
43
+ subprocess.run("pip install ray[tune]".split(), check=True) # do not add single quotes here
44
44
 
45
45
  import ray
46
46
  from ray import tune
@@ -48,7 +48,7 @@ def run_ray_tune(
48
48
  from ray.air.integrations.wandb import WandbLoggerCallback
49
49
  from ray.tune.schedulers import ASHAScheduler
50
50
  except ImportError:
51
- raise ModuleNotFoundError('Ray Tune required but not found. To install run: pip install "ray[tune]<=2.9.3"')
51
+ raise ModuleNotFoundError('Ray Tune required but not found. To install run: pip install "ray[tune]"')
52
52
 
53
53
  try:
54
54
  import wandb
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: ultralytics
3
- Version: 8.1.37
3
+ Version: 8.1.39
4
4
  Summary: Ultralytics YOLOv8 for SOTA object detection, multi-object tracking, instance segmentation, pose estimation and image classification.
5
5
  Author: Glenn Jocher, Ayush Chaurasia, Jing Qiu
6
6
  Maintainer: Glenn Jocher, Ayush Chaurasia, Jing Qiu