ultralytics 8.3.189__py3-none-any.whl → 8.3.191__py3-none-any.whl

ultralytics/data/utils.py

@@ -1,5 +1,7 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
+from __future__ import annotations
+
 import json
 import os
 import random
@@ -9,7 +11,7 @@ import zipfile
 from multiprocessing.pool import ThreadPool
 from pathlib import Path
 from tarfile import is_tarfile
-from typing import Any, Dict, List, Tuple, Union
+from typing import Any
 
 import cv2
 import numpy as np
@@ -39,14 +41,14 @@ VID_FORMATS = {"asf", "avi", "gif", "m4v", "mkv", "mov", "mp4", "mpeg", "mpg", "
 FORMATS_HELP_MSG = f"Supported formats are:\nimages: {IMG_FORMATS}\nvideos: {VID_FORMATS}"
 
 
-def img2label_paths(img_paths: List[str]) -> List[str]:
+def img2label_paths(img_paths: list[str]) -> list[str]:
     """Convert image paths to label paths by replacing 'images' with 'labels' and extension with '.txt'."""
     sa, sb = f"{os.sep}images{os.sep}", f"{os.sep}labels{os.sep}"  # /images/, /labels/ substrings
     return [sb.join(x.rsplit(sa, 1)).rsplit(".", 1)[0] + ".txt" for x in img_paths]
 
 
 def check_file_speeds(
-    files: List[str], threshold_ms: float = 10, threshold_mb: float = 50, max_files: int = 5, prefix: str = ""
+    files: list[str], threshold_ms: float = 10, threshold_mb: float = 50, max_files: int = 5, prefix: str = ""
 ):
     """
     Check dataset file access speed and provide performance feedback.
@@ -66,7 +68,7 @@ def check_file_speeds(
         >>> image_files = list(Path("dataset/images").glob("*.jpg"))
         >>> check_file_speeds(image_files, threshold_ms=15)
     """
-    if not files or len(files) == 0:
+    if not files:
         LOGGER.warning(f"{prefix}Image speed checks: No files to check")
         return
 
@@ -123,7 +125,7 @@ def check_file_speeds(
         )
 
 
-def get_hash(paths: List[str]) -> str:
+def get_hash(paths: list[str]) -> str:
     """Return a single hash value of a list of paths (files or dirs)."""
     size = 0
     for p in paths:
@@ -136,7 +138,7 @@ def get_hash(paths: List[str]) -> str:
     return h.hexdigest()  # return hash
 
 
-def exif_size(img: Image.Image) -> Tuple[int, int]:
+def exif_size(img: Image.Image) -> tuple[int, int]:
     """Return exif-corrected PIL size."""
     s = img.size  # (width, height)
     if img.format == "JPEG":  # only support JPEG images
@@ -150,7 +152,7 @@ def exif_size(img: Image.Image) -> Tuple[int, int]:
     return s
 
 
-def verify_image(args: Tuple) -> Tuple:
+def verify_image(args: tuple) -> tuple:
     """Verify one image."""
     (im_file, cls), prefix = args
     # Number (found, corrupt), message
@@ -175,7 +177,7 @@ def verify_image(args: Tuple) -> Tuple:
     return (im_file, cls), nf, nc, msg
 
 
-def verify_image_label(args: Tuple) -> List:
+def verify_image_label(args: tuple) -> list:
     """Verify one image-label pair."""
     im_file, lb_file, prefix, keypoint, num_cls, nkpt, ndim, single_cls = args
     # Number (missing, found, empty, corrupt), message, segments, keypoints
@@ -247,7 +249,7 @@ def verify_image_label(args: Tuple) -> List:
         return [None, None, None, None, None, nm, nf, ne, nc, msg]
 
 
-def visualize_image_annotations(image_path: str, txt_path: str, label_map: Dict[int, str]):
+def visualize_image_annotations(image_path: str, txt_path: str, label_map: dict[int, str]):
     """
     Visualize YOLO annotations (bounding boxes and class labels) on an image.
 
@@ -292,7 +294,7 @@ def visualize_image_annotations(image_path: str, txt_path: str, label_map: Dict[
 
 
 def polygon2mask(
-    imgsz: Tuple[int, int], polygons: List[np.ndarray], color: int = 1, downsample_ratio: int = 1
+    imgsz: tuple[int, int], polygons: list[np.ndarray], color: int = 1, downsample_ratio: int = 1
 ) -> np.ndarray:
     """
     Convert a list of polygons to a binary mask of the specified image size.
@@ -317,7 +319,7 @@ def polygon2mask(
 
 
 def polygons2masks(
-    imgsz: Tuple[int, int], polygons: List[np.ndarray], color: int, downsample_ratio: int = 1
+    imgsz: tuple[int, int], polygons: list[np.ndarray], color: int, downsample_ratio: int = 1
 ) -> np.ndarray:
     """
     Convert a list of polygons to a set of binary masks of the specified image size.
@@ -336,8 +338,8 @@ def polygons2masks(
 
 
 def polygons2masks_overlap(
-    imgsz: Tuple[int, int], segments: List[np.ndarray], downsample_ratio: int = 1
-) -> Tuple[np.ndarray, np.ndarray]:
+    imgsz: tuple[int, int], segments: list[np.ndarray], downsample_ratio: int = 1
+) -> tuple[np.ndarray, np.ndarray]:
     """Return a (640, 640) overlap mask."""
     masks = np.zeros(
         (imgsz[0] // downsample_ratio, imgsz[1] // downsample_ratio),
@@ -345,8 +347,13 @@ def polygons2masks_overlap(
     )
     areas = []
     ms = []
-    for si in range(len(segments)):
-        mask = polygon2mask(imgsz, [segments[si].reshape(-1)], downsample_ratio=downsample_ratio, color=1)
+    for segment in segments:
+        mask = polygon2mask(
+            imgsz,
+            [segment.reshape(-1)],
+            downsample_ratio=downsample_ratio,
+            color=1,
+        )
         ms.append(mask.astype(masks.dtype))
         areas.append(mask.sum())
     areas = np.asarray(areas)
@@ -380,7 +387,7 @@ def find_dataset_yaml(path: Path) -> Path:
     return files[0]
 
 
-def check_det_dataset(dataset: str, autodownload: bool = True) -> Dict[str, Any]:
+def check_det_dataset(dataset: str, autodownload: bool = True) -> dict[str, Any]:
     """
     Download, verify, and/or unzip a dataset if not found locally.
 
@@ -464,7 +471,7 @@ def check_det_dataset(dataset: str, autodownload: bool = True) -> Dict[str, Any]
                 safe_download(url=s, dir=DATASETS_DIR, delete=True)
             elif s.startswith("bash "):  # bash script
                 LOGGER.info(f"Running {s} ...")
-                r = os.system(s)
+                subprocess.run(s.split(), check=True)
             else:  # python script
                 exec(s, {"yaml": data})
             dt = f"({round(time.time() - t, 1)}s)"
@@ -475,7 +482,7 @@ def check_det_dataset(dataset: str, autodownload: bool = True) -> Dict[str, Any]
     return data  # dictionary
 
 
-def check_cls_dataset(dataset: Union[str, Path], split: str = "") -> Dict[str, Any]:
+def check_cls_dataset(dataset: str | Path, split: str = "") -> dict[str, Any]:
     """
     Check a classification dataset such as Imagenet.
 
@@ -509,7 +516,7 @@ def check_cls_dataset(dataset: Union[str, Path], split: str = "") -> Dict[str, A
         LOGGER.warning(f"Dataset not found, missing path {data_dir}, attempting download...")
         t = time.time()
         if str(dataset) == "imagenet":
-            subprocess.run(f"bash {ROOT / 'data/scripts/get_imagenet.sh'}", shell=True, check=True)
+            subprocess.run(["bash", str(ROOT / "data/scripts/get_imagenet.sh")], check=True)
         else:
             url = f"https://github.com/ultralytics/assets/releases/download/v0.0.0/{dataset}.zip"
             download(url, dir=data_dir.parent)
@@ -517,8 +524,7 @@ def check_cls_dataset(dataset: Union[str, Path], split: str = "") -> Dict[str, A
     train_set = data_dir / "train"
     if not train_set.is_dir():
         LOGGER.warning(f"Dataset 'split=train' not found at {train_set}")
-        image_files = list(data_dir.rglob("*.jpg")) + list(data_dir.rglob("*.png"))
-        if image_files:
+        if image_files := list(data_dir.rglob("*.jpg")) + list(data_dir.rglob("*.png")):
             from ultralytics.data.split import split_classify_dataset
 
             LOGGER.info(f"Found {len(image_files)} images in subdirectories. Attempting to split...")
@@ -632,7 +638,7 @@ class HUBDatasetStats:
         self.data = data
 
     @staticmethod
-    def _unzip(path: Path) -> Tuple[bool, str, Path]:
+    def _unzip(path: Path) -> tuple[bool, str, Path]:
         """Unzip data.zip."""
         if not str(path).endswith(".zip"):  # path is data.yaml
             return False, None, path
@@ -646,7 +652,7 @@ class HUBDatasetStats:
         """Save a compressed image for HUB previews."""
         compress_one_image(f, self.im_dir / Path(f).name)  # save to dataset-hub
 
-    def get_json(self, save: bool = False, verbose: bool = False) -> Dict:
+    def get_json(self, save: bool = False, verbose: bool = False) -> dict:
         """Return dataset JSON for Ultralytics HUB."""
 
         def _round(labels):
@@ -773,7 +779,7 @@ def compress_one_image(f: str, f_new: str = None, max_dim: int = 1920, quality:
         cv2.imwrite(str(f_new or f), im)
 
 
-def load_dataset_cache_file(path: Path) -> Dict:
+def load_dataset_cache_file(path: Path) -> dict:
     """Load an Ultralytics *.cache dictionary from path."""
     import gc
 
@@ -783,7 +789,7 @@ def load_dataset_cache_file(path: Path) -> Dict:
     return cache
 
 
-def save_dataset_cache_file(prefix: str, path: Path, x: Dict, version: str):
+def save_dataset_cache_file(prefix: str, path: Path, x: dict, version: str):
     """Save an Ultralytics dataset *.cache dictionary x to path."""
     x["version"] = version  # add cache version
     if is_dir_writeable(path.parent):

ultralytics/engine/exporter.py

@@ -107,7 +107,9 @@ from ultralytics.utils.checks import (
 from ultralytics.utils.downloads import attempt_download_asset, get_github_assets, safe_download
 from ultralytics.utils.export import export_engine, export_onnx
 from ultralytics.utils.files import file_size, spaces_in_path
-from ultralytics.utils.ops import Profile, nms_rotated
+from ultralytics.utils.metrics import batch_probiou
+from ultralytics.utils.nms import TorchNMS
+from ultralytics.utils.ops import Profile
 from ultralytics.utils.patches import arange_patch
 from ultralytics.utils.torch_utils import TORCH_1_13, get_latest_opset, select_device
 
@@ -347,7 +349,7 @@ class Exporter:
             assert not getattr(model, "end2end", False), "TFLite INT8 export not supported for end2end models."
         if self.args.nms:
             assert not isinstance(model, ClassificationModel), "'nms=True' is not valid for classification models."
-            assert not (tflite and ARM64 and LINUX), "TFLite export with NMS unsupported on ARM64 Linux"
+            assert not tflite or not ARM64 or not LINUX, "TFLite export with NMS unsupported on ARM64 Linux"
             if getattr(model, "end2end", False):
                 LOGGER.warning("'nms=True' is not available for end2end models. Forcing 'nms=False'.")
                 self.args.nms = False
@@ -434,7 +436,7 @@ class Exporter:
 
         y = None
         for _ in range(2):  # dry runs
-            y = NMSModel(model, self.args)(im) if self.args.nms and not (coreml or imx) else model(im)
+            y = NMSModel(model, self.args)(im) if self.args.nms and not coreml and not imx else model(im)
         if self.args.half and onnx and self.device.type != "cpu":
             im, model = im.half(), model.half()  # to FP16
 
@@ -1562,12 +1564,13 @@ class NMSModel(torch.nn.Module):
                 nmsbox = torch.cat((offbox, nmsbox[:, end:]), dim=-1)
             nms_fn = (
                 partial(
-                    nms_rotated,
+                    TorchNMS.fast_nms,
                     use_triu=not (
                         self.is_tf
                         or (self.args.opset or 14) < 14
                         or (self.args.format == "openvino" and self.args.int8)  # OpenVINO int8 error with triu
                     ),
+                    iou_func=batch_probiou,
                 )
                 if self.obb
                 else nms

ultralytics/engine/model.py

@@ -1,8 +1,10 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
+from __future__ import annotations
+
 import inspect
 from pathlib import Path
-from typing import Any, Dict, List, Union
+from typing import Any
 
 import numpy as np
 import torch
@@ -79,7 +81,7 @@ class Model(torch.nn.Module):
 
     def __init__(
         self,
-        model: Union[str, Path, "Model"] = "yolo11n.pt",
+        model: str | Path | Model = "yolo11n.pt",
         task: str = None,
         verbose: bool = False,
     ) -> None:
@@ -155,7 +157,7 @@ class Model(torch.nn.Module):
 
     def __call__(
         self,
-        source: Union[str, Path, int, Image.Image, list, tuple, np.ndarray, torch.Tensor] = None,
+        source: str | Path | int | Image.Image | list | tuple | np.ndarray | torch.Tensor = None,
         stream: bool = False,
         **kwargs: Any,
     ) -> list:
@@ -333,7 +335,7 @@ class Model(torch.nn.Module):
                 f"argument directly in your inference command, i.e. 'model.predict(source=..., device=0)'"
             )
 
-    def reset_weights(self) -> "Model":
+    def reset_weights(self) -> Model:
         """
         Reset the model's weights to their initial state.
 
@@ -359,7 +361,7 @@ class Model(torch.nn.Module):
             p.requires_grad = True
         return self
 
-    def load(self, weights: Union[str, Path] = "yolo11n.pt") -> "Model":
+    def load(self, weights: str | Path = "yolo11n.pt") -> Model:
         """
         Load parameters from the specified weights file into the model.
 
@@ -387,7 +389,7 @@ class Model(torch.nn.Module):
         self.model.load(weights)
         return self
 
-    def save(self, filename: Union[str, Path] = "saved_model.pt") -> None:
+    def save(self, filename: str | Path = "saved_model.pt") -> None:
         """
         Save the current model state to a file.
 
@@ -464,7 +466,7 @@ class Model(torch.nn.Module):
 
     def embed(
         self,
-        source: Union[str, Path, int, list, tuple, np.ndarray, torch.Tensor] = None,
+        source: str | Path | int | list | tuple | np.ndarray | torch.Tensor = None,
         stream: bool = False,
         **kwargs: Any,
     ) -> list:
@@ -495,11 +497,11 @@ class Model(torch.nn.Module):
 
     def predict(
         self,
-        source: Union[str, Path, int, Image.Image, list, tuple, np.ndarray, torch.Tensor] = None,
+        source: str | Path | int | Image.Image | list | tuple | np.ndarray | torch.Tensor = None,
         stream: bool = False,
         predictor=None,
         **kwargs: Any,
-    ) -> List[Results]:
+    ) -> list[Results]:
         """
         Perform predictions on the given image source using the YOLO model.
 
@@ -556,11 +558,11 @@ class Model(torch.nn.Module):
 
     def track(
         self,
-        source: Union[str, Path, int, list, tuple, np.ndarray, torch.Tensor] = None,
+        source: str | Path | int | list | tuple | np.ndarray | torch.Tensor = None,
         stream: bool = False,
         persist: bool = False,
         **kwargs: Any,
-    ) -> List[Results]:
+    ) -> list[Results]:
         """
         Conduct object tracking on the specified input source using the registered trackers.
 
@@ -853,7 +855,7 @@ class Model(torch.nn.Module):
             args = {**self.overrides, **custom, **kwargs, "mode": "train"}  # highest priority args on the right
             return Tuner(args=args, _callbacks=self.callbacks)(model=self, iterations=iterations)
 
-    def _apply(self, fn) -> "Model":
+    def _apply(self, fn) -> Model:
         """
         Apply a function to model tensors that are not parameters or registered buffers.
 
@@ -882,7 +884,7 @@ class Model(torch.nn.Module):
         return self
 
     @property
-    def names(self) -> Dict[int, str]:
+    def names(self) -> dict[int, str]:
         """
         Retrieve the class names associated with the loaded model.
 
@@ -1036,7 +1038,7 @@ class Model(torch.nn.Module):
             self.callbacks[event] = [callbacks.default_callbacks[event][0]]
 
     @staticmethod
-    def _reset_ckpt_args(args: Dict[str, Any]) -> Dict[str, Any]:
+    def _reset_ckpt_args(args: dict[str, Any]) -> dict[str, Any]:
         """
         Reset specific arguments when loading a PyTorch model checkpoint.
 

ultralytics/engine/predictor.py

@@ -32,11 +32,13 @@ Usage - formats:
                               yolo11n_rknn_model         # Rockchip RKNN
 """
 
+from __future__ import annotations
+
 import platform
 import re
 import threading
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Union
+from typing import Any
 
 import cv2
 import numpy as np
@@ -109,8 +111,8 @@ class BasePredictor:
     def __init__(
         self,
         cfg=DEFAULT_CFG,
-        overrides: Optional[Dict[str, Any]] = None,
-        _callbacks: Optional[Dict[str, List[callable]]] = None,
+        overrides: dict[str, Any] | None = None,
+        _callbacks: dict[str, list[callable]] | None = None,
     ):
         """
         Initialize the BasePredictor class.
@@ -147,7 +149,7 @@ class BasePredictor:
         self._lock = threading.Lock()  # for automatic thread-safe inference
         callbacks.add_integration_callbacks(self)
 
-    def preprocess(self, im: Union[torch.Tensor, List[np.ndarray]]) -> torch.Tensor:
+    def preprocess(self, im: torch.Tensor | list[np.ndarray]) -> torch.Tensor:
         """
         Prepare input image before inference.
 
@@ -181,7 +183,7 @@ class BasePredictor:
         )
         return self.model(im, augment=self.args.augment, visualize=visualize, embed=self.args.embed, *args, **kwargs)
 
-    def pre_transform(self, im: List[np.ndarray]) -> List[np.ndarray]:
+    def pre_transform(self, im: list[np.ndarray]) -> list[np.ndarray]:
         """
         Pre-transform input image before inference.
 
@@ -389,7 +391,7 @@ class BasePredictor:
             verbose (bool): Whether to print verbose output.
         """
         self.model = AutoBackend(
-            weights=model or self.args.model,
+            model=model or self.args.model,
             device=select_device(self.args.device, verbose=verbose),
             dnn=self.args.dnn,
             data=self.args.data,
@@ -404,7 +406,7 @@ class BasePredictor:
             self.args.imgsz = self.model.imgsz  # reuse imgsz from export metadata
         self.model.eval()
 
-    def write_results(self, i: int, p: Path, im: torch.Tensor, s: List[str]) -> str:
+    def write_results(self, i: int, p: Path, im: torch.Tensor, s: list[str]) -> str:
         """
         Write inference results to a file or directory.