dgenerate-ultralytics-headless 8.3.190__py3-none-any.whl → 8.3.192__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.
- {dgenerate_ultralytics_headless-8.3.190.dist-info → dgenerate_ultralytics_headless-8.3.192.dist-info}/METADATA +1 -1
- {dgenerate_ultralytics_headless-8.3.190.dist-info → dgenerate_ultralytics_headless-8.3.192.dist-info}/RECORD +103 -102
- tests/test_cuda.py +6 -5
- tests/test_exports.py +1 -6
- tests/test_python.py +1 -4
- tests/test_solutions.py +1 -1
- ultralytics/__init__.py +1 -1
- ultralytics/cfg/__init__.py +16 -14
- ultralytics/cfg/datasets/SKU-110K.yaml +1 -1
- ultralytics/cfg/datasets/VisDrone.yaml +4 -4
- ultralytics/data/annotator.py +6 -6
- ultralytics/data/augment.py +53 -51
- ultralytics/data/base.py +15 -13
- ultralytics/data/build.py +7 -4
- ultralytics/data/converter.py +9 -10
- ultralytics/data/dataset.py +24 -22
- ultralytics/data/loaders.py +13 -11
- ultralytics/data/split.py +4 -3
- ultralytics/data/split_dota.py +14 -12
- ultralytics/data/utils.py +29 -23
- ultralytics/engine/exporter.py +2 -2
- ultralytics/engine/model.py +16 -14
- ultralytics/engine/predictor.py +8 -6
- ultralytics/engine/results.py +54 -52
- ultralytics/engine/trainer.py +8 -3
- ultralytics/engine/tuner.py +230 -42
- ultralytics/hub/google/__init__.py +7 -6
- ultralytics/hub/session.py +8 -6
- ultralytics/hub/utils.py +3 -4
- ultralytics/models/fastsam/model.py +8 -6
- ultralytics/models/nas/model.py +5 -3
- ultralytics/models/rtdetr/train.py +4 -3
- ultralytics/models/rtdetr/val.py +6 -4
- ultralytics/models/sam/amg.py +13 -10
- ultralytics/models/sam/model.py +3 -2
- ultralytics/models/sam/modules/blocks.py +21 -21
- ultralytics/models/sam/modules/decoders.py +11 -11
- ultralytics/models/sam/modules/encoders.py +25 -25
- ultralytics/models/sam/modules/memory_attention.py +9 -8
- ultralytics/models/sam/modules/sam.py +8 -10
- ultralytics/models/sam/modules/tiny_encoder.py +21 -20
- ultralytics/models/sam/modules/transformer.py +6 -5
- ultralytics/models/sam/modules/utils.py +7 -5
- ultralytics/models/sam/predict.py +32 -31
- ultralytics/models/utils/loss.py +29 -27
- ultralytics/models/utils/ops.py +10 -8
- ultralytics/models/yolo/classify/train.py +9 -7
- ultralytics/models/yolo/classify/val.py +11 -9
- ultralytics/models/yolo/detect/predict.py +1 -1
- ultralytics/models/yolo/detect/train.py +8 -6
- ultralytics/models/yolo/detect/val.py +22 -20
- ultralytics/models/yolo/model.py +14 -14
- ultralytics/models/yolo/obb/train.py +5 -3
- ultralytics/models/yolo/obb/val.py +11 -9
- ultralytics/models/yolo/pose/train.py +7 -5
- ultralytics/models/yolo/pose/val.py +12 -10
- ultralytics/models/yolo/segment/train.py +4 -5
- ultralytics/models/yolo/segment/val.py +13 -11
- ultralytics/models/yolo/world/train.py +10 -8
- ultralytics/models/yolo/yoloe/train.py +10 -10
- ultralytics/models/yolo/yoloe/val.py +11 -9
- ultralytics/nn/autobackend.py +17 -19
- ultralytics/nn/modules/block.py +12 -12
- ultralytics/nn/modules/conv.py +4 -3
- ultralytics/nn/modules/head.py +41 -37
- ultralytics/nn/modules/transformer.py +22 -21
- ultralytics/nn/tasks.py +2 -2
- ultralytics/nn/text_model.py +6 -5
- ultralytics/solutions/analytics.py +7 -5
- ultralytics/solutions/config.py +12 -10
- ultralytics/solutions/distance_calculation.py +3 -3
- ultralytics/solutions/heatmap.py +4 -2
- ultralytics/solutions/object_counter.py +5 -3
- ultralytics/solutions/parking_management.py +4 -2
- ultralytics/solutions/region_counter.py +7 -5
- ultralytics/solutions/similarity_search.py +5 -3
- ultralytics/solutions/solutions.py +38 -36
- ultralytics/solutions/streamlit_inference.py +8 -7
- ultralytics/trackers/bot_sort.py +11 -9
- ultralytics/trackers/byte_tracker.py +17 -15
- ultralytics/trackers/utils/gmc.py +4 -3
- ultralytics/utils/__init__.py +16 -88
- ultralytics/utils/autobatch.py +3 -2
- ultralytics/utils/autodevice.py +10 -10
- ultralytics/utils/benchmarks.py +11 -10
- ultralytics/utils/callbacks/comet.py +9 -9
- ultralytics/utils/checks.py +17 -26
- ultralytics/utils/export.py +12 -11
- ultralytics/utils/files.py +8 -7
- ultralytics/utils/git.py +139 -0
- ultralytics/utils/instance.py +8 -7
- ultralytics/utils/loss.py +15 -13
- ultralytics/utils/metrics.py +62 -62
- ultralytics/utils/ops.py +3 -2
- ultralytics/utils/patches.py +6 -4
- ultralytics/utils/plotting.py +20 -18
- ultralytics/utils/torch_utils.py +4 -2
- ultralytics/utils/tqdm.py +18 -14
- ultralytics/utils/triton.py +3 -2
- {dgenerate_ultralytics_headless-8.3.190.dist-info → dgenerate_ultralytics_headless-8.3.192.dist-info}/WHEEL +0 -0
- {dgenerate_ultralytics_headless-8.3.190.dist-info → dgenerate_ultralytics_headless-8.3.192.dist-info}/entry_points.txt +0 -0
- {dgenerate_ultralytics_headless-8.3.190.dist-info → dgenerate_ultralytics_headless-8.3.192.dist-info}/licenses/LICENSE +0 -0
- {dgenerate_ultralytics_headless-8.3.190.dist-info → dgenerate_ultralytics_headless-8.3.192.dist-info}/top_level.txt +0 -0
@@ -1,9 +1,11 @@
|
|
1
1
|
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
2
2
|
|
3
|
+
from __future__ import annotations
|
4
|
+
|
3
5
|
import math
|
4
6
|
import random
|
5
7
|
from copy import copy
|
6
|
-
from typing import Any
|
8
|
+
from typing import Any
|
7
9
|
|
8
10
|
import numpy as np
|
9
11
|
import torch.nn as nn
|
@@ -51,7 +53,7 @@ class DetectionTrainer(BaseTrainer):
|
|
51
53
|
>>> trainer.train()
|
52
54
|
"""
|
53
55
|
|
54
|
-
def build_dataset(self, img_path: str, mode: str = "train", batch:
|
56
|
+
def build_dataset(self, img_path: str, mode: str = "train", batch: int | None = None):
|
55
57
|
"""
|
56
58
|
Build YOLO Dataset for training or validation.
|
57
59
|
|
@@ -89,7 +91,7 @@ class DetectionTrainer(BaseTrainer):
|
|
89
91
|
workers = self.args.workers if mode == "train" else self.args.workers * 2
|
90
92
|
return build_dataloader(dataset, batch_size, workers, shuffle, rank) # return dataloader
|
91
93
|
|
92
|
-
def preprocess_batch(self, batch:
|
94
|
+
def preprocess_batch(self, batch: dict) -> dict:
|
93
95
|
"""
|
94
96
|
Preprocess a batch of images by scaling and converting to float.
|
95
97
|
|
@@ -127,7 +129,7 @@ class DetectionTrainer(BaseTrainer):
|
|
127
129
|
self.model.args = self.args # attach hyperparameters to model
|
128
130
|
# TODO: self.model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc
|
129
131
|
|
130
|
-
def get_model(self, cfg:
|
132
|
+
def get_model(self, cfg: str | None = None, weights: str | None = None, verbose: bool = True):
|
131
133
|
"""
|
132
134
|
Return a YOLO detection model.
|
133
135
|
|
@@ -151,7 +153,7 @@ class DetectionTrainer(BaseTrainer):
|
|
151
153
|
self.test_loader, save_dir=self.save_dir, args=copy(self.args), _callbacks=self.callbacks
|
152
154
|
)
|
153
155
|
|
154
|
-
def label_loss_items(self, loss_items:
|
156
|
+
def label_loss_items(self, loss_items: list[float] | None = None, prefix: str = "train"):
|
155
157
|
"""
|
156
158
|
Return a loss dict with labeled training loss items tensor.
|
157
159
|
|
@@ -179,7 +181,7 @@ class DetectionTrainer(BaseTrainer):
|
|
179
181
|
"Size",
|
180
182
|
)
|
181
183
|
|
182
|
-
def plot_training_samples(self, batch:
|
184
|
+
def plot_training_samples(self, batch: dict[str, Any], ni: int) -> None:
|
183
185
|
"""
|
184
186
|
Plot training samples with their annotations.
|
185
187
|
|
@@ -1,8 +1,10 @@
|
|
1
1
|
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
2
2
|
|
3
|
+
from __future__ import annotations
|
4
|
+
|
3
5
|
import os
|
4
6
|
from pathlib import Path
|
5
|
-
from typing import Any
|
7
|
+
from typing import Any
|
6
8
|
|
7
9
|
import numpy as np
|
8
10
|
import torch
|
@@ -59,7 +61,7 @@ class DetectionValidator(BaseValidator):
|
|
59
61
|
self.niou = self.iouv.numel()
|
60
62
|
self.metrics = DetMetrics()
|
61
63
|
|
62
|
-
def preprocess(self, batch:
|
64
|
+
def preprocess(self, batch: dict[str, Any]) -> dict[str, Any]:
|
63
65
|
"""
|
64
66
|
Preprocess batch of images for YOLO validation.
|
65
67
|
|
@@ -72,7 +74,7 @@ class DetectionValidator(BaseValidator):
|
|
72
74
|
batch["img"] = batch["img"].to(self.device, non_blocking=True)
|
73
75
|
batch["img"] = (batch["img"].half() if self.args.half else batch["img"].float()) / 255
|
74
76
|
for k in {"batch_idx", "cls", "bboxes"}:
|
75
|
-
batch[k] = batch[k].to(self.device)
|
77
|
+
batch[k] = batch[k].to(self.device, non_blocking=True)
|
76
78
|
|
77
79
|
return batch
|
78
80
|
|
@@ -104,7 +106,7 @@ class DetectionValidator(BaseValidator):
|
|
104
106
|
"""Return a formatted string summarizing class metrics of YOLO model."""
|
105
107
|
return ("%22s" + "%11s" * 6) % ("Class", "Images", "Instances", "Box(P", "R", "mAP50", "mAP50-95)")
|
106
108
|
|
107
|
-
def postprocess(self, preds: torch.Tensor) ->
|
109
|
+
def postprocess(self, preds: torch.Tensor) -> list[dict[str, torch.Tensor]]:
|
108
110
|
"""
|
109
111
|
Apply Non-maximum suppression to prediction outputs.
|
110
112
|
|
@@ -128,7 +130,7 @@ class DetectionValidator(BaseValidator):
|
|
128
130
|
)
|
129
131
|
return [{"bboxes": x[:, :4], "conf": x[:, 4], "cls": x[:, 5], "extra": x[:, 6:]} for x in outputs]
|
130
132
|
|
131
|
-
def _prepare_batch(self, si: int, batch:
|
133
|
+
def _prepare_batch(self, si: int, batch: dict[str, Any]) -> dict[str, Any]:
|
132
134
|
"""
|
133
135
|
Prepare a batch of images and annotations for validation.
|
134
136
|
|
@@ -156,7 +158,7 @@ class DetectionValidator(BaseValidator):
|
|
156
158
|
"im_file": batch["im_file"][si],
|
157
159
|
}
|
158
160
|
|
159
|
-
def _prepare_pred(self, pred:
|
161
|
+
def _prepare_pred(self, pred: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
|
160
162
|
"""
|
161
163
|
Prepare predictions for evaluation against ground truth.
|
162
164
|
|
@@ -170,7 +172,7 @@ class DetectionValidator(BaseValidator):
|
|
170
172
|
pred["cls"] *= 0
|
171
173
|
return pred
|
172
174
|
|
173
|
-
def update_metrics(self, preds:
|
175
|
+
def update_metrics(self, preds: list[dict[str, torch.Tensor]], batch: dict[str, Any]) -> None:
|
174
176
|
"""
|
175
177
|
Update metrics with new predictions and ground truth.
|
176
178
|
|
@@ -225,7 +227,7 @@ class DetectionValidator(BaseValidator):
|
|
225
227
|
self.metrics.confusion_matrix = self.confusion_matrix
|
226
228
|
self.metrics.save_dir = self.save_dir
|
227
229
|
|
228
|
-
def get_stats(self) ->
|
230
|
+
def get_stats(self) -> dict[str, Any]:
|
229
231
|
"""
|
230
232
|
Calculate and return metrics statistics.
|
231
233
|
|
@@ -256,7 +258,7 @@ class DetectionValidator(BaseValidator):
|
|
256
258
|
)
|
257
259
|
)
|
258
260
|
|
259
|
-
def _process_batch(self, preds:
|
261
|
+
def _process_batch(self, preds: dict[str, torch.Tensor], batch: dict[str, Any]) -> dict[str, np.ndarray]:
|
260
262
|
"""
|
261
263
|
Return correct prediction matrix.
|
262
264
|
|
@@ -272,7 +274,7 @@ class DetectionValidator(BaseValidator):
|
|
272
274
|
iou = box_iou(batch["bboxes"], preds["bboxes"])
|
273
275
|
return {"tp": self.match_predictions(preds["cls"], batch["cls"], iou).cpu().numpy()}
|
274
276
|
|
275
|
-
def build_dataset(self, img_path: str, mode: str = "val", batch:
|
277
|
+
def build_dataset(self, img_path: str, mode: str = "val", batch: int | None = None) -> torch.utils.data.Dataset:
|
276
278
|
"""
|
277
279
|
Build YOLO Dataset.
|
278
280
|
|
@@ -300,7 +302,7 @@ class DetectionValidator(BaseValidator):
|
|
300
302
|
dataset = self.build_dataset(dataset_path, batch=batch_size, mode="val")
|
301
303
|
return build_dataloader(dataset, batch_size, self.args.workers, shuffle=False, rank=-1) # return dataloader
|
302
304
|
|
303
|
-
def plot_val_samples(self, batch:
|
305
|
+
def plot_val_samples(self, batch: dict[str, Any], ni: int) -> None:
|
304
306
|
"""
|
305
307
|
Plot validation image samples.
|
306
308
|
|
@@ -317,7 +319,7 @@ class DetectionValidator(BaseValidator):
|
|
317
319
|
)
|
318
320
|
|
319
321
|
def plot_predictions(
|
320
|
-
self, batch:
|
322
|
+
self, batch: dict[str, Any], preds: list[dict[str, torch.Tensor]], ni: int, max_det: int | None = None
|
321
323
|
) -> None:
|
322
324
|
"""
|
323
325
|
Plot predicted bounding boxes on input images and save the result.
|
@@ -345,7 +347,7 @@ class DetectionValidator(BaseValidator):
|
|
345
347
|
on_plot=self.on_plot,
|
346
348
|
) # pred
|
347
349
|
|
348
|
-
def save_one_txt(self, predn:
|
350
|
+
def save_one_txt(self, predn: dict[str, torch.Tensor], save_conf: bool, shape: tuple[int, int], file: Path) -> None:
|
349
351
|
"""
|
350
352
|
Save YOLO detections to a txt file in normalized coordinates in a specific format.
|
351
353
|
|
@@ -364,7 +366,7 @@ class DetectionValidator(BaseValidator):
|
|
364
366
|
boxes=torch.cat([predn["bboxes"], predn["conf"].unsqueeze(-1), predn["cls"].unsqueeze(-1)], dim=1),
|
365
367
|
).save_txt(file, save_conf=save_conf)
|
366
368
|
|
367
|
-
def pred_to_json(self, predn:
|
369
|
+
def pred_to_json(self, predn: dict[str, torch.Tensor], pbatch: dict[str, Any]) -> None:
|
368
370
|
"""
|
369
371
|
Serialize YOLO predictions to COCO json format.
|
370
372
|
|
@@ -398,7 +400,7 @@ class DetectionValidator(BaseValidator):
|
|
398
400
|
}
|
399
401
|
)
|
400
402
|
|
401
|
-
def scale_preds(self, predn:
|
403
|
+
def scale_preds(self, predn: dict[str, torch.Tensor], pbatch: dict[str, Any]) -> dict[str, torch.Tensor]:
|
402
404
|
"""Scales predictions to the original image size."""
|
403
405
|
return {
|
404
406
|
**predn,
|
@@ -410,7 +412,7 @@ class DetectionValidator(BaseValidator):
|
|
410
412
|
),
|
411
413
|
}
|
412
414
|
|
413
|
-
def eval_json(self, stats:
|
415
|
+
def eval_json(self, stats: dict[str, Any]) -> dict[str, Any]:
|
414
416
|
"""
|
415
417
|
Evaluate YOLO output in JSON format and return performance statistics.
|
416
418
|
|
@@ -430,12 +432,12 @@ class DetectionValidator(BaseValidator):
|
|
430
432
|
|
431
433
|
def coco_evaluate(
|
432
434
|
self,
|
433
|
-
stats:
|
435
|
+
stats: dict[str, Any],
|
434
436
|
pred_json: str,
|
435
437
|
anno_json: str,
|
436
|
-
iou_types:
|
437
|
-
suffix:
|
438
|
-
) ->
|
438
|
+
iou_types: str | list[str] = "bbox",
|
439
|
+
suffix: str | list[str] = "Box",
|
440
|
+
) -> dict[str, Any]:
|
439
441
|
"""
|
440
442
|
Evaluate COCO/LVIS metrics using faster-coco-eval library.
|
441
443
|
|
ultralytics/models/yolo/model.py
CHANGED
@@ -1,7 +1,9 @@
|
|
1
1
|
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
2
2
|
|
3
|
+
from __future__ import annotations
|
4
|
+
|
3
5
|
from pathlib import Path
|
4
|
-
from typing import Any
|
6
|
+
from typing import Any
|
5
7
|
|
6
8
|
import torch
|
7
9
|
|
@@ -49,7 +51,7 @@ class YOLO(Model):
|
|
49
51
|
>>> model = YOLO("yolo11n.yaml")
|
50
52
|
"""
|
51
53
|
|
52
|
-
def __init__(self, model:
|
54
|
+
def __init__(self, model: str | Path = "yolo11n.pt", task: str | None = None, verbose: bool = False):
|
53
55
|
"""
|
54
56
|
Initialize a YOLO model.
|
55
57
|
|
@@ -87,7 +89,7 @@ class YOLO(Model):
|
|
87
89
|
self.__dict__ = new_instance.__dict__
|
88
90
|
|
89
91
|
@property
|
90
|
-
def task_map(self) ->
|
92
|
+
def task_map(self) -> dict[str, dict[str, Any]]:
|
91
93
|
"""Map head to model, trainer, validator, and predictor classes."""
|
92
94
|
return {
|
93
95
|
"classify": {
|
@@ -149,7 +151,7 @@ class YOLOWorld(Model):
|
|
149
151
|
>>> model.set_classes(["person", "car", "bicycle"])
|
150
152
|
"""
|
151
153
|
|
152
|
-
def __init__(self, model:
|
154
|
+
def __init__(self, model: str | Path = "yolov8s-world.pt", verbose: bool = False) -> None:
|
153
155
|
"""
|
154
156
|
Initialize YOLOv8-World model with a pre-trained model file.
|
155
157
|
|
@@ -167,7 +169,7 @@ class YOLOWorld(Model):
|
|
167
169
|
self.model.names = YAML.load(ROOT / "cfg/datasets/coco8.yaml").get("names")
|
168
170
|
|
169
171
|
@property
|
170
|
-
def task_map(self) ->
|
172
|
+
def task_map(self) -> dict[str, dict[str, Any]]:
|
171
173
|
"""Map head to model, validator, and predictor classes."""
|
172
174
|
return {
|
173
175
|
"detect": {
|
@@ -178,7 +180,7 @@ class YOLOWorld(Model):
|
|
178
180
|
}
|
179
181
|
}
|
180
182
|
|
181
|
-
def set_classes(self, classes:
|
183
|
+
def set_classes(self, classes: list[str]) -> None:
|
182
184
|
"""
|
183
185
|
Set the model's class names for detection.
|
184
186
|
|
@@ -232,9 +234,7 @@ class YOLOE(Model):
|
|
232
234
|
>>> results = model.predict("image.jpg", visual_prompts=prompts)
|
233
235
|
"""
|
234
236
|
|
235
|
-
def __init__(
|
236
|
-
self, model: Union[str, Path] = "yoloe-11s-seg.pt", task: Optional[str] = None, verbose: bool = False
|
237
|
-
) -> None:
|
237
|
+
def __init__(self, model: str | Path = "yoloe-11s-seg.pt", task: str | None = None, verbose: bool = False) -> None:
|
238
238
|
"""
|
239
239
|
Initialize YOLOE model with a pre-trained model file.
|
240
240
|
|
@@ -246,7 +246,7 @@ class YOLOE(Model):
|
|
246
246
|
super().__init__(model=model, task=task, verbose=verbose)
|
247
247
|
|
248
248
|
@property
|
249
|
-
def task_map(self) ->
|
249
|
+
def task_map(self) -> dict[str, dict[str, Any]]:
|
250
250
|
"""Map head to model, validator, and predictor classes."""
|
251
251
|
return {
|
252
252
|
"detect": {
|
@@ -291,7 +291,7 @@ class YOLOE(Model):
|
|
291
291
|
assert isinstance(self.model, YOLOEModel)
|
292
292
|
return self.model.get_visual_pe(img, visual)
|
293
293
|
|
294
|
-
def set_vocab(self, vocab:
|
294
|
+
def set_vocab(self, vocab: list[str], names: list[str]) -> None:
|
295
295
|
"""
|
296
296
|
Set vocabulary and class names for the YOLOE model.
|
297
297
|
|
@@ -317,7 +317,7 @@ class YOLOE(Model):
|
|
317
317
|
assert isinstance(self.model, YOLOEModel)
|
318
318
|
return self.model.get_vocab(names)
|
319
319
|
|
320
|
-
def set_classes(self, classes:
|
320
|
+
def set_classes(self, classes: list[str], embeddings: torch.Tensor | None = None) -> None:
|
321
321
|
"""
|
322
322
|
Set the model's class names and embeddings for detection.
|
323
323
|
|
@@ -341,7 +341,7 @@ class YOLOE(Model):
|
|
341
341
|
self,
|
342
342
|
validator=None,
|
343
343
|
load_vp: bool = False,
|
344
|
-
refer_data:
|
344
|
+
refer_data: str | None = None,
|
345
345
|
**kwargs,
|
346
346
|
):
|
347
347
|
"""
|
@@ -368,7 +368,7 @@ class YOLOE(Model):
|
|
368
368
|
self,
|
369
369
|
source=None,
|
370
370
|
stream: bool = False,
|
371
|
-
visual_prompts:
|
371
|
+
visual_prompts: dict[str, list] = {},
|
372
372
|
refer_image=None,
|
373
373
|
predictor=yolo.yoloe.YOLOEVPDetectPredictor,
|
374
374
|
**kwargs,
|
@@ -1,8 +1,10 @@
|
|
1
1
|
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
2
2
|
|
3
|
+
from __future__ import annotations
|
4
|
+
|
3
5
|
from copy import copy
|
4
6
|
from pathlib import Path
|
5
|
-
from typing import Any
|
7
|
+
from typing import Any
|
6
8
|
|
7
9
|
from ultralytics.models import yolo
|
8
10
|
from ultralytics.nn.tasks import OBBModel
|
@@ -31,7 +33,7 @@ class OBBTrainer(yolo.detect.DetectionTrainer):
|
|
31
33
|
>>> trainer.train()
|
32
34
|
"""
|
33
35
|
|
34
|
-
def __init__(self, cfg=DEFAULT_CFG, overrides:
|
36
|
+
def __init__(self, cfg=DEFAULT_CFG, overrides: dict | None = None, _callbacks: list[Any] | None = None):
|
35
37
|
"""
|
36
38
|
Initialize an OBBTrainer object for training Oriented Bounding Box (OBB) models.
|
37
39
|
|
@@ -57,7 +59,7 @@ class OBBTrainer(yolo.detect.DetectionTrainer):
|
|
57
59
|
super().__init__(cfg, overrides, _callbacks)
|
58
60
|
|
59
61
|
def get_model(
|
60
|
-
self, cfg:
|
62
|
+
self, cfg: str | dict | None = None, weights: str | Path | None = None, verbose: bool = True
|
61
63
|
) -> OBBModel:
|
62
64
|
"""
|
63
65
|
Return OBBModel initialized with specified config and weights.
|
@@ -1,7 +1,9 @@
|
|
1
1
|
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
2
2
|
|
3
|
+
from __future__ import annotations
|
4
|
+
|
3
5
|
from pathlib import Path
|
4
|
-
from typing import Any
|
6
|
+
from typing import Any
|
5
7
|
|
6
8
|
import numpy as np
|
7
9
|
import torch
|
@@ -70,7 +72,7 @@ class OBBValidator(DetectionValidator):
|
|
70
72
|
self.is_dota = isinstance(val, str) and "DOTA" in val # check if dataset is DOTA format
|
71
73
|
self.confusion_matrix.task = "obb" # set confusion matrix task to 'obb'
|
72
74
|
|
73
|
-
def _process_batch(self, preds:
|
75
|
+
def _process_batch(self, preds: dict[str, torch.Tensor], batch: dict[str, torch.Tensor]) -> dict[str, np.ndarray]:
|
74
76
|
"""
|
75
77
|
Compute the correct prediction matrix for a batch of detections and ground truth bounding boxes.
|
76
78
|
|
@@ -96,7 +98,7 @@ class OBBValidator(DetectionValidator):
|
|
96
98
|
iou = batch_probiou(batch["bboxes"], preds["bboxes"])
|
97
99
|
return {"tp": self.match_predictions(preds["cls"], batch["cls"], iou).cpu().numpy()}
|
98
100
|
|
99
|
-
def postprocess(self, preds: torch.Tensor) ->
|
101
|
+
def postprocess(self, preds: torch.Tensor) -> list[dict[str, torch.Tensor]]:
|
100
102
|
"""
|
101
103
|
Args:
|
102
104
|
preds (torch.Tensor): Raw predictions from the model.
|
@@ -109,7 +111,7 @@ class OBBValidator(DetectionValidator):
|
|
109
111
|
pred["bboxes"] = torch.cat([pred["bboxes"], pred.pop("extra")], dim=-1) # concatenate angle
|
110
112
|
return preds
|
111
113
|
|
112
|
-
def _prepare_batch(self, si: int, batch:
|
114
|
+
def _prepare_batch(self, si: int, batch: dict[str, Any]) -> dict[str, Any]:
|
113
115
|
"""
|
114
116
|
Prepare batch data for OBB validation with proper scaling and formatting.
|
115
117
|
|
@@ -143,7 +145,7 @@ class OBBValidator(DetectionValidator):
|
|
143
145
|
"im_file": batch["im_file"][si],
|
144
146
|
}
|
145
147
|
|
146
|
-
def plot_predictions(self, batch:
|
148
|
+
def plot_predictions(self, batch: dict[str, Any], preds: list[torch.Tensor], ni: int) -> None:
|
147
149
|
"""
|
148
150
|
Plot predicted bounding boxes on input images and save the result.
|
149
151
|
|
@@ -163,7 +165,7 @@ class OBBValidator(DetectionValidator):
|
|
163
165
|
p["bboxes"][:, :4] = ops.xywh2xyxy(p["bboxes"][:, :4]) # convert to xyxy format for plotting
|
164
166
|
super().plot_predictions(batch, preds, ni) # plot bboxes
|
165
167
|
|
166
|
-
def pred_to_json(self, predn:
|
168
|
+
def pred_to_json(self, predn: dict[str, torch.Tensor], pbatch: dict[str, Any]) -> None:
|
167
169
|
"""
|
168
170
|
Convert YOLO predictions to COCO JSON format with rotated bounding box information.
|
169
171
|
|
@@ -194,7 +196,7 @@ class OBBValidator(DetectionValidator):
|
|
194
196
|
}
|
195
197
|
)
|
196
198
|
|
197
|
-
def save_one_txt(self, predn:
|
199
|
+
def save_one_txt(self, predn: dict[str, torch.Tensor], save_conf: bool, shape: tuple[int, int], file: Path) -> None:
|
198
200
|
"""
|
199
201
|
Save YOLO OBB detections to a text file in normalized coordinates.
|
200
202
|
|
@@ -221,7 +223,7 @@ class OBBValidator(DetectionValidator):
|
|
221
223
|
obb=torch.cat([predn["bboxes"], predn["conf"].unsqueeze(-1), predn["cls"].unsqueeze(-1)], dim=1),
|
222
224
|
).save_txt(file, save_conf=save_conf)
|
223
225
|
|
224
|
-
def scale_preds(self, predn:
|
226
|
+
def scale_preds(self, predn: dict[str, torch.Tensor], pbatch: dict[str, Any]) -> dict[str, torch.Tensor]:
|
225
227
|
"""Scales predictions to the original image size."""
|
226
228
|
return {
|
227
229
|
**predn,
|
@@ -230,7 +232,7 @@ class OBBValidator(DetectionValidator):
|
|
230
232
|
),
|
231
233
|
}
|
232
234
|
|
233
|
-
def eval_json(self, stats:
|
235
|
+
def eval_json(self, stats: dict[str, Any]) -> dict[str, Any]:
|
234
236
|
"""
|
235
237
|
Evaluate YOLO output in JSON format and save predictions in DOTA format.
|
236
238
|
|
@@ -1,8 +1,10 @@
|
|
1
1
|
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
2
2
|
|
3
|
+
from __future__ import annotations
|
4
|
+
|
3
5
|
from copy import copy
|
4
6
|
from pathlib import Path
|
5
|
-
from typing import Any
|
7
|
+
from typing import Any
|
6
8
|
|
7
9
|
from ultralytics.models import yolo
|
8
10
|
from ultralytics.nn.tasks import PoseModel
|
@@ -38,7 +40,7 @@ class PoseTrainer(yolo.detect.DetectionTrainer):
|
|
38
40
|
>>> trainer.train()
|
39
41
|
"""
|
40
42
|
|
41
|
-
def __init__(self, cfg=DEFAULT_CFG, overrides:
|
43
|
+
def __init__(self, cfg=DEFAULT_CFG, overrides: dict[str, Any] | None = None, _callbacks=None):
|
42
44
|
"""
|
43
45
|
Initialize a PoseTrainer object for training YOLO pose estimation models.
|
44
46
|
|
@@ -73,8 +75,8 @@ class PoseTrainer(yolo.detect.DetectionTrainer):
|
|
73
75
|
|
74
76
|
def get_model(
|
75
77
|
self,
|
76
|
-
cfg:
|
77
|
-
weights:
|
78
|
+
cfg: str | Path | dict[str, Any] | None = None,
|
79
|
+
weights: str | Path | None = None,
|
78
80
|
verbose: bool = True,
|
79
81
|
) -> PoseModel:
|
80
82
|
"""
|
@@ -112,7 +114,7 @@ class PoseTrainer(yolo.detect.DetectionTrainer):
|
|
112
114
|
"""Plot training/validation metrics."""
|
113
115
|
plot_results(file=self.csv, pose=True, on_plot=self.on_plot) # save results.png
|
114
116
|
|
115
|
-
def get_dataset(self) ->
|
117
|
+
def get_dataset(self) -> dict[str, Any]:
|
116
118
|
"""
|
117
119
|
Retrieve the dataset and ensure it contains the required `kpt_shape` key.
|
118
120
|
|
@@ -1,7 +1,9 @@
|
|
1
1
|
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
2
2
|
|
3
|
+
from __future__ import annotations
|
4
|
+
|
3
5
|
from pathlib import Path
|
4
|
-
from typing import Any
|
6
|
+
from typing import Any
|
5
7
|
|
6
8
|
import numpy as np
|
7
9
|
import torch
|
@@ -81,10 +83,10 @@ class PoseValidator(DetectionValidator):
|
|
81
83
|
"See https://github.com/ultralytics/ultralytics/issues/4031."
|
82
84
|
)
|
83
85
|
|
84
|
-
def preprocess(self, batch:
|
86
|
+
def preprocess(self, batch: dict[str, Any]) -> dict[str, Any]:
|
85
87
|
"""Preprocess batch by converting keypoints data to float and moving it to the device."""
|
86
88
|
batch = super().preprocess(batch)
|
87
|
-
batch["keypoints"] = batch["keypoints"].to(self.device).float()
|
89
|
+
batch["keypoints"] = batch["keypoints"].to(self.device, non_blocking=True).float()
|
88
90
|
return batch
|
89
91
|
|
90
92
|
def get_desc(self) -> str:
|
@@ -116,7 +118,7 @@ class PoseValidator(DetectionValidator):
|
|
116
118
|
nkpt = self.kpt_shape[0]
|
117
119
|
self.sigma = OKS_SIGMA if is_pose else np.ones(nkpt) / nkpt
|
118
120
|
|
119
|
-
def postprocess(self, preds: torch.Tensor) ->
|
121
|
+
def postprocess(self, preds: torch.Tensor) -> dict[str, torch.Tensor]:
|
120
122
|
"""
|
121
123
|
Postprocess YOLO predictions to extract and reshape keypoints for pose estimation.
|
122
124
|
|
@@ -146,7 +148,7 @@ class PoseValidator(DetectionValidator):
|
|
146
148
|
pred["keypoints"] = pred.pop("extra").view(-1, *self.kpt_shape) # remove extra if exists
|
147
149
|
return preds
|
148
150
|
|
149
|
-
def _prepare_batch(self, si: int, batch:
|
151
|
+
def _prepare_batch(self, si: int, batch: dict[str, Any]) -> dict[str, Any]:
|
150
152
|
"""
|
151
153
|
Prepare a batch for processing by converting keypoints to float and scaling to original dimensions.
|
152
154
|
|
@@ -170,7 +172,7 @@ class PoseValidator(DetectionValidator):
|
|
170
172
|
pbatch["keypoints"] = kpts
|
171
173
|
return pbatch
|
172
174
|
|
173
|
-
def _process_batch(self, preds:
|
175
|
+
def _process_batch(self, preds: dict[str, torch.Tensor], batch: dict[str, Any]) -> dict[str, np.ndarray]:
|
174
176
|
"""
|
175
177
|
Return correct prediction matrix by computing Intersection over Union (IoU) between detections and ground truth.
|
176
178
|
|
@@ -200,7 +202,7 @@ class PoseValidator(DetectionValidator):
|
|
200
202
|
tp.update({"tp_p": tp_p}) # update tp with kpts IoU
|
201
203
|
return tp
|
202
204
|
|
203
|
-
def save_one_txt(self, predn:
|
205
|
+
def save_one_txt(self, predn: dict[str, torch.Tensor], save_conf: bool, shape: tuple[int, int], file: Path) -> None:
|
204
206
|
"""
|
205
207
|
Save YOLO pose detections to a text file in normalized coordinates.
|
206
208
|
|
@@ -224,7 +226,7 @@ class PoseValidator(DetectionValidator):
|
|
224
226
|
keypoints=predn["keypoints"],
|
225
227
|
).save_txt(file, save_conf=save_conf)
|
226
228
|
|
227
|
-
def pred_to_json(self, predn:
|
229
|
+
def pred_to_json(self, predn: dict[str, torch.Tensor], pbatch: dict[str, Any]) -> None:
|
228
230
|
"""
|
229
231
|
Convert YOLO predictions to COCO JSON format.
|
230
232
|
|
@@ -246,7 +248,7 @@ class PoseValidator(DetectionValidator):
|
|
246
248
|
for i, k in enumerate(kpts.flatten(1, 2).tolist()):
|
247
249
|
self.jdict[-len(kpts) + i]["keypoints"] = k # keypoints
|
248
250
|
|
249
|
-
def scale_preds(self, predn:
|
251
|
+
def scale_preds(self, predn: dict[str, torch.Tensor], pbatch: dict[str, Any]) -> dict[str, torch.Tensor]:
|
250
252
|
"""Scales predictions to the original image size."""
|
251
253
|
return {
|
252
254
|
**super().scale_preds(predn, pbatch),
|
@@ -258,7 +260,7 @@ class PoseValidator(DetectionValidator):
|
|
258
260
|
),
|
259
261
|
}
|
260
262
|
|
261
|
-
def eval_json(self, stats:
|
263
|
+
def eval_json(self, stats: dict[str, Any]) -> dict[str, Any]:
|
262
264
|
"""Evaluate object detection model using COCO JSON format."""
|
263
265
|
anno_json = self.data["path"] / "annotations/person_keypoints_val2017.json" # annotations
|
264
266
|
pred_json = self.save_dir / "predictions.json" # predictions
|
@@ -1,8 +1,9 @@
|
|
1
1
|
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
2
2
|
|
3
|
+
from __future__ import annotations
|
4
|
+
|
3
5
|
from copy import copy
|
4
6
|
from pathlib import Path
|
5
|
-
from typing import Dict, Optional, Union
|
6
7
|
|
7
8
|
from ultralytics.models import yolo
|
8
9
|
from ultralytics.nn.tasks import SegmentationModel
|
@@ -27,7 +28,7 @@ class SegmentationTrainer(yolo.detect.DetectionTrainer):
|
|
27
28
|
>>> trainer.train()
|
28
29
|
"""
|
29
30
|
|
30
|
-
def __init__(self, cfg=DEFAULT_CFG, overrides:
|
31
|
+
def __init__(self, cfg=DEFAULT_CFG, overrides: dict | None = None, _callbacks=None):
|
31
32
|
"""
|
32
33
|
Initialize a SegmentationTrainer object.
|
33
34
|
|
@@ -50,9 +51,7 @@ class SegmentationTrainer(yolo.detect.DetectionTrainer):
|
|
50
51
|
overrides["task"] = "segment"
|
51
52
|
super().__init__(cfg, overrides, _callbacks)
|
52
53
|
|
53
|
-
def get_model(
|
54
|
-
self, cfg: Optional[Union[Dict, str]] = None, weights: Optional[Union[str, Path]] = None, verbose: bool = True
|
55
|
-
):
|
54
|
+
def get_model(self, cfg: dict | str | None = None, weights: str | Path | None = None, verbose: bool = True):
|
56
55
|
"""
|
57
56
|
Initialize and return a SegmentationModel with specified configuration and weights.
|
58
57
|
|