dgenerate-ultralytics-headless 8.3.214__py3-none-any.whl → 8.3.248__py3-none-any.whl

ultralytics/utils/export/imx.py

@@ -3,20 +3,51 @@
 from __future__ import annotations
 
 import subprocess
+import sys
 import types
 from pathlib import Path
+from shutil import which
 
+import numpy as np
 import torch
 
-from ultralytics.nn.modules import Detect, Pose
-from ultralytics.utils import LOGGER
+from ultralytics.nn.modules import Detect, Pose, Segment
+from ultralytics.utils import LOGGER, WINDOWS
+from ultralytics.utils.patches import onnx_export_patch
 from ultralytics.utils.tal import make_anchors
 from ultralytics.utils.torch_utils import copy_attr
 
+# Configuration for Model Compression Toolkit (MCT) quantization
+MCT_CONFIG = {
+    "YOLO11": {
+        "detect": {
+            "layer_names": ["sub", "mul_2", "add_14", "cat_21"],
+            "weights_memory": 2585350.2439,
+            "n_layers": 238,
+        },
+        "pose": {
+            "layer_names": ["sub", "mul_2", "add_14", "cat_22", "cat_23", "mul_4", "add_15"],
+            "weights_memory": 2437771.67,
+            "n_layers": 257,
+        },
+        "classify": {"layer_names": [], "weights_memory": np.inf, "n_layers": 112},
+        "segment": {"layer_names": ["sub", "mul_2", "add_14", "cat_22"], "weights_memory": 2466604.8, "n_layers": 265},
+    },
+    "YOLOv8": {
+        "detect": {"layer_names": ["sub", "mul", "add_6", "cat_17"], "weights_memory": 2550540.8, "n_layers": 168},
+        "pose": {
+            "layer_names": ["add_7", "mul_2", "cat_19", "mul", "sub", "add_6", "cat_18"],
+            "weights_memory": 2482451.85,
+            "n_layers": 187,
+        },
+        "classify": {"layer_names": [], "weights_memory": np.inf, "n_layers": 73},
+        "segment": {"layer_names": ["sub", "mul", "add_6", "cat_18"], "weights_memory": 2580060.0, "n_layers": 195},
+    },
+}
+
 
 class FXModel(torch.nn.Module):
-    """
-    A custom model class for torch.fx compatibility.
+    """A custom model class for torch.fx compatibility.
 
     This class extends `torch.nn.Module` and is designed to ensure compatibility with torch.fx for tracing and graph
     manipulation. It copies attributes from an existing model and explicitly sets the model attribute to ensure proper
@@ -27,8 +58,7 @@ class FXModel(torch.nn.Module):
     """
 
     def __init__(self, model, imgsz=(640, 640)):
-        """
-        Initialize the FXModel.
+        """Initialize the FXModel.
 
         Args:
             model (nn.Module): The original model to wrap for torch.fx compatibility.
@@ -41,8 +71,7 @@ class FXModel(torch.nn.Module):
         self.imgsz = imgsz
 
     def forward(self, x):
-        """
-        Forward pass through the model.
+        """Forward pass through the model.
 
         This method performs the forward pass through the model, handling the dependencies between layers and saving
         intermediate outputs.
@@ -68,6 +97,8 @@ class FXModel(torch.nn.Module):
                 )
             if type(m) is Pose:
                 m.forward = types.MethodType(pose_forward, m)  # bind method to Detect
+            if type(m) is Segment:
+                m.forward = types.MethodType(segment_forward, m)  # bind method to Detect
             x = m(x)  # run
             y.append(x)  # save output
         return x
@@ -87,11 +118,20 @@ def pose_forward(self, x: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tenso
     kpt = torch.cat([self.cv4[i](x[i]).view(bs, self.nk, -1) for i in range(self.nl)], -1)  # (bs, 17*3, h*w)
     x = Detect.forward(self, x)
     pred_kpt = self.kpts_decode(bs, kpt)
-    return (*x, pred_kpt.permute(0, 2, 1))
+    return *x, pred_kpt.permute(0, 2, 1)
+
+
+def segment_forward(self, x: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Forward pass for imx segmentation."""
+    p = self.proto(x[0])  # mask protos
+    bs = p.shape[0]  # batch size
+    mc = torch.cat([self.cv4[i](x[i]).view(bs, self.nm, -1) for i in range(self.nl)], 2)  # mask coefficients
+    x = Detect.forward(self, x)
+    return *x, mc.transpose(1, 2), p
 
 
 class NMSWrapper(torch.nn.Module):
-    """Wrap PyTorch Module with multiclass_nms layer from sony_custom_layers."""
+    """Wrap PyTorch Module with multiclass_nms layer from edge-mdt-cl."""
 
     def __init__(
         self,
@@ -101,8 +141,7 @@ class NMSWrapper(torch.nn.Module):
         max_detections: int = 300,
         task: str = "detect",
     ):
-        """
-        Initialize NMSWrapper with PyTorch Module and NMS parameters.
+        """Initialize NMSWrapper with PyTorch Module and NMS parameters.
 
         Args:
             model (torch.nn.Module): Model instance.
@@ -120,7 +159,7 @@ class NMSWrapper(torch.nn.Module):
 
     def forward(self, images):
         """Forward pass with model inference and NMS post-processing."""
-        from sony_custom_layers.pytorch import multiclass_nms_with_indices
+        from edgemdt_cl.pytorch.nms.nms_with_indices import multiclass_nms_with_indices
 
         # model inference
         outputs = self.model(images)
@@ -136,6 +175,10 @@ class NMSWrapper(torch.nn.Module):
             kpts = outputs[2]  # (bs, max_detections, kpts 17*3)
             out_kpts = torch.gather(kpts, 1, nms_outputs.indices.unsqueeze(-1).expand(-1, -1, kpts.size(-1)))
             return nms_outputs.boxes, nms_outputs.scores, nms_outputs.labels, out_kpts
+        if self.task == "segment":
+            mc, proto = outputs[2], outputs[3]
+            out_mc = torch.gather(mc, 1, nms_outputs.indices.unsqueeze(-1).expand(-1, -1, mc.size(-1)))
+            return nms_outputs.boxes, nms_outputs.scores, nms_outputs.labels, out_mc, proto
         return nms_outputs.boxes, nms_outputs.scores, nms_outputs.labels, nms_outputs.n_valid
 
 
@@ -150,12 +193,10 @@ def torch2imx(
     dataset=None,
     prefix: str = "",
 ):
-    """
-    Export YOLO model to IMX format for deployment on Sony IMX500 devices.
+    """Export YOLO model to IMX format for deployment on Sony IMX500 devices.
 
-    This function quantizes a YOLO model using Model Compression Toolkit (MCT) and exports it
-    to IMX format compatible with Sony IMX500 edge devices. It supports both YOLOv8n and YOLO11n
-    models for detection and pose estimation tasks.
+    This function quantizes a YOLO model using Model Compression Toolkit (MCT) and exports it to IMX format compatible
+    with Sony IMX500 edge devices. It supports both YOLOv8n and YOLO11n models for detection and pose estimation tasks.
 
     Args:
         model (torch.nn.Module): The YOLO model to export. Must be YOLOv8n or YOLO11n.
@@ -164,8 +205,8 @@ def torch2imx(
         iou (float): IoU threshold for NMS post-processing.
         max_det (int): Maximum number of detections to return.
         metadata (dict | None, optional): Metadata to embed in the ONNX model. Defaults to None.
-        gptq (bool, optional): Whether to use Gradient-Based Post Training Quantization.
-            If False, uses standard Post Training Quantization. Defaults to False.
+        gptq (bool, optional): Whether to use Gradient-Based Post Training Quantization. If False, uses standard Post
+            Training Quantization. Defaults to False.
         dataset (optional): Representative dataset for quantization calibration. Defaults to None.
         prefix (str, optional): Logging prefix string. Defaults to "".
 
@@ -175,13 +216,13 @@ def torch2imx(
     Raises:
         ValueError: If the model is not a supported YOLOv8n or YOLO11n variant.
 
-    Example:
+    Examples:
         >>> from ultralytics import YOLO
         >>> model = YOLO("yolo11n.pt")
         >>> path, _ = export_imx(model, "model.imx", conf=0.25, iou=0.45, max_det=300)
 
-    Note:
-        - Requires model_compression_toolkit, onnx, edgemdt_tpc, and sony_custom_layers packages
+    Notes:
+        - Requires model_compression_toolkit, onnx, edgemdt_tpc, and edge-mdt-cl packages
         - Only supports YOLOv8n and YOLO11n models (detection and pose tasks)
         - Output includes quantized ONNX model, IMX binary, and labels.txt file
     """
@@ -197,33 +238,17 @@ def torch2imx(
             img = img / 255.0
             yield [img]
 
+    # NOTE: need tpc_version to be "4.0" for IMX500 Pose estimation models
     tpc = get_target_platform_capabilities(tpc_version="4.0", device_type="imx500")
 
     bit_cfg = mct.core.BitWidthConfig()
-    if "C2PSA" in model.__str__():  # YOLO11
-        if model.task == "detect":
-            layer_names = ["sub", "mul_2", "add_14", "cat_21"]
-            weights_memory = 2585350.2439
-            n_layers = 238  # 238 layers for fused YOLO11n
-        elif model.task == "pose":
-            layer_names = ["sub", "mul_2", "add_14", "cat_22", "cat_23", "mul_4", "add_15"]
-            weights_memory = 2437771.67
-            n_layers = 257  # 257 layers for fused YOLO11n-pose
-    else:  # YOLOv8
-        if model.task == "detect":
-            layer_names = ["sub", "mul", "add_6", "cat_17"]
-            weights_memory = 2550540.8
-            n_layers = 168  # 168 layers for fused YOLOv8n
-        elif model.task == "pose":
-            layer_names = ["add_7", "mul_2", "cat_19", "mul", "sub", "add_6", "cat_18"]
-            weights_memory = 2482451.85
-            n_layers = 187  # 187 layers for fused YOLO11n-pose
+    mct_config = MCT_CONFIG["YOLO11" if "C2PSA" in model.__str__() else "YOLOv8"][model.task]
 
     # Check if the model has the expected number of layers
-    if len(list(model.modules())) != n_layers:
+    if len(list(model.modules())) != mct_config["n_layers"]:
         raise ValueError("IMX export only supported for YOLOv8n and YOLO11n models.")
 
-    for layer_name in layer_names:
+    for layer_name in mct_config["layer_names"]:
         bit_cfg.set_manual_activation_bit_width([mct.core.common.network_editors.NodeNameFilter(layer_name)], 16)
 
     config = mct.core.CoreConfig(
@@ -232,7 +257,7 @@ def torch2imx(
         bit_width_config=bit_cfg,
     )
 
-    resource_utilization = mct.core.ResourceUtilization(weights_memory=weights_memory)
+    resource_utilization = mct.core.ResourceUtilization(weights_memory=mct_config["weights_memory"])
 
     quant_model = (
         mct.gptq.pytorch_gradient_post_training_quantization(  # Perform Gradient-Based Post Training Quantization
@@ -255,20 +280,23 @@ def torch2imx(
         )[0]
     )
 
-    quant_model = NMSWrapper(
-        model=quant_model,
-        score_threshold=conf or 0.001,
-        iou_threshold=iou,
-        max_detections=max_det,
-        task=model.task,
-    )
+    if model.task != "classify":
+        quant_model = NMSWrapper(
+            model=quant_model,
+            score_threshold=conf or 0.001,
+            iou_threshold=iou,
+            max_detections=max_det,
+            task=model.task,
+        )
 
     f = Path(str(file).replace(file.suffix, "_imx_model"))
     f.mkdir(exist_ok=True)
     onnx_model = f / Path(str(file.name).replace(file.suffix, "_imx.onnx"))  # js dir
-    mct.exporter.pytorch_export_model(
-        model=quant_model, save_model_path=onnx_model, repr_dataset=representative_dataset_gen
-    )
+
+    with onnx_export_patch():
+        mct.exporter.pytorch_export_model(
+            model=quant_model, save_model_path=onnx_model, repr_dataset=representative_dataset_gen
+        )
 
     model_onnx = onnx.load(onnx_model)  # load onnx model
     for k, v in metadata.items():
@@ -277,8 +305,16 @@ def torch2imx(
 
     onnx.save(model_onnx, onnx_model)
 
+    # Find imxconv-pt binary - check venv bin directory first, then PATH
+    bin_dir = Path(sys.executable).parent
+    imxconv = bin_dir / ("imxconv-pt.exe" if WINDOWS else "imxconv-pt")
+    if not imxconv.exists():
+        imxconv = which("imxconv-pt")  # fallback to PATH
+    if not imxconv:
+        raise FileNotFoundError("imxconv-pt not found. Install with: pip install imx500-converter[pt]")
+
     subprocess.run(
-        ["imxconv-pt", "-i", str(onnx_model), "-o", str(f), "--no-input-persistency", "--overwrite-output"],
+        [str(imxconv), "-i", str(onnx_model), "-o", str(f), "--no-input-persistency", "--overwrite-output"],
         check=True,
     )
 

ultralytics/utils/export/tensorflow.py

@@ -0,0 +1,231 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from ultralytics.nn.modules import Detect, Pose
+from ultralytics.utils import LOGGER
+from ultralytics.utils.downloads import attempt_download_asset
+from ultralytics.utils.files import spaces_in_path
+from ultralytics.utils.tal import make_anchors
+
+
+def tf_wrapper(model: torch.nn.Module) -> torch.nn.Module:
+    """A wrapper to add TensorFlow compatible inference methods to Detect and Pose layers."""
+    for m in model.modules():
+        if not isinstance(m, Detect):
+            continue
+        import types
+
+        m._inference = types.MethodType(_tf_inference, m)
+        if type(m) is Pose:
+            m.kpts_decode = types.MethodType(tf_kpts_decode, m)
+    return model
+
+
+def _tf_inference(self, x: list[torch.Tensor]) -> tuple[torch.Tensor]:
+    """Decode boxes and cls scores for tf object detection."""
+    shape = x[0].shape  # BCHW
+    x_cat = torch.cat([xi.view(x[0].shape[0], self.no, -1) for xi in x], 2)
+    box, cls = x_cat.split((self.reg_max * 4, self.nc), 1)
+    if self.dynamic or self.shape != shape:
+        self.anchors, self.strides = (x.transpose(0, 1) for x in make_anchors(x, self.stride, 0.5))
+        self.shape = shape
+    grid_h, grid_w = shape[2], shape[3]
+    grid_size = torch.tensor([grid_w, grid_h, grid_w, grid_h], device=box.device).reshape(1, 4, 1)
+    norm = self.strides / (self.stride[0] * grid_size)
+    dbox = self.decode_bboxes(self.dfl(box) * norm, self.anchors.unsqueeze(0) * norm[:, :2])
+    return torch.cat((dbox, cls.sigmoid()), 1)
+
+
+def tf_kpts_decode(self, bs: int, kpts: torch.Tensor) -> torch.Tensor:
+    """Decode keypoints for tf pose estimation."""
+    ndim = self.kpt_shape[1]
+    # required for TFLite export to avoid 'PLACEHOLDER_FOR_GREATER_OP_CODES' bug
+    # Precompute normalization factor to increase numerical stability
+    y = kpts.view(bs, *self.kpt_shape, -1)
+    grid_h, grid_w = self.shape[2], self.shape[3]
+    grid_size = torch.tensor([grid_w, grid_h], device=y.device).reshape(1, 2, 1)
+    norm = self.strides / (self.stride[0] * grid_size)
+    a = (y[:, :, :2] * 2.0 + (self.anchors - 0.5)) * norm
+    if ndim == 3:
+        a = torch.cat((a, y[:, :, 2:3].sigmoid()), 2)
+    return a.view(bs, self.nk, -1)
+
+
+def onnx2saved_model(
+    onnx_file: str,
+    output_dir: Path,
+    int8: bool = False,
+    images: np.ndarray = None,
+    disable_group_convolution: bool = False,
+    prefix="",
+):
+    """Convert a ONNX model to TensorFlow SavedModel format via ONNX.
+
+    Args:
+        onnx_file (str): ONNX file path.
+        output_dir (Path): Output directory path for the SavedModel.
+        int8 (bool, optional): Enable INT8 quantization. Defaults to False.
+        images (np.ndarray, optional): Calibration images for INT8 quantization in BHWC format.
+        disable_group_convolution (bool, optional): Disable group convolution optimization. Defaults to False.
+        prefix (str, optional): Logging prefix. Defaults to "".
+
+    Returns:
+        (keras.Model): Converted Keras model.
+
+    Notes:
+        - Requires onnx2tf package. Downloads calibration data if INT8 quantization is enabled.
+        - Removes temporary files and renames quantized models after conversion.
+    """
+    # Pre-download calibration file to fix https://github.com/PINTO0309/onnx2tf/issues/545
+    onnx2tf_file = Path("calibration_image_sample_data_20x128x128x3_float32.npy")
+    if not onnx2tf_file.exists():
+        attempt_download_asset(f"{onnx2tf_file}.zip", unzip=True, delete=True)
+    np_data = None
+    if int8:
+        tmp_file = output_dir / "tmp_tflite_int8_calibration_images.npy"  # int8 calibration images file
+        if images is not None:
+            output_dir.mkdir(parents=True, exist_ok=True)
+            np.save(str(tmp_file), images)  # BHWC
+            np_data = [["images", tmp_file, [[[[0, 0, 0]]]], [[[[255, 255, 255]]]]]]
+
+    # Patch onnx.helper for onnx_graphsurgeon compatibility with ONNX>=1.17
+    # The float32_to_bfloat16 function was removed in ONNX 1.17, but onnx_graphsurgeon still uses it
+    import onnx.helper
+
+    if not hasattr(onnx.helper, "float32_to_bfloat16"):
+        import struct
+
+        def float32_to_bfloat16(fval):
+            """Convert float32 to bfloat16 (truncates lower 16 bits of mantissa)."""
+            ival = struct.unpack("=I", struct.pack("=f", fval))[0]
+            return ival >> 16
+
+        onnx.helper.float32_to_bfloat16 = float32_to_bfloat16
+
+    import onnx2tf  # scoped for after ONNX export for reduced conflict during import
+
+    LOGGER.info(f"{prefix} starting TFLite export with onnx2tf {onnx2tf.__version__}...")
+    keras_model = onnx2tf.convert(
+        input_onnx_file_path=onnx_file,
+        output_folder_path=str(output_dir),
+        not_use_onnxsim=True,
+        verbosity="error",  # note INT8-FP16 activation bug https://github.com/ultralytics/ultralytics/issues/15873
+        output_integer_quantized_tflite=int8,
+        custom_input_op_name_np_data_path=np_data,
+        enable_batchmatmul_unfold=True and not int8,  # fix lower no. of detected objects on GPU delegate
+        output_signaturedefs=True,  # fix error with Attention block group convolution
+        disable_group_convolution=disable_group_convolution,  # fix error with group convolution
+    )
+
+    # Remove/rename TFLite models
+    if int8:
+        tmp_file.unlink(missing_ok=True)
+        for file in output_dir.rglob("*_dynamic_range_quant.tflite"):
+            file.rename(file.with_name(file.stem.replace("_dynamic_range_quant", "_int8") + file.suffix))
+        for file in output_dir.rglob("*_integer_quant_with_int16_act.tflite"):
+            file.unlink()  # delete extra fp16 activation TFLite files
+    return keras_model
+
+
+def keras2pb(keras_model, file: Path, prefix=""):
+    """Convert a Keras model to TensorFlow GraphDef (.pb) format.
+
+    Args:
+        keras_model (keras.Model): Keras model to convert to frozen graph format.
+        file (Path): Output file path (suffix will be changed to .pb).
+        prefix (str, optional): Logging prefix. Defaults to "".
+
+    Notes:
+        Creates a frozen graph by converting variables to constants for inference optimization.
+    """
+    import tensorflow as tf
+    from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
+
+    LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...")
+    m = tf.function(lambda x: keras_model(x))  # full model
+    m = m.get_concrete_function(tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype))
+    frozen_func = convert_variables_to_constants_v2(m)
+    frozen_func.graph.as_graph_def()
+    tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(file.parent), name=file.name, as_text=False)
+
+
+def tflite2edgetpu(tflite_file: str | Path, output_dir: str | Path, prefix: str = ""):
+    """Convert a TensorFlow Lite model to Edge TPU format using the Edge TPU compiler.
+
+    Args:
+        tflite_file (str | Path): Path to the input TensorFlow Lite (.tflite) model file.
+        output_dir (str | Path): Output directory path for the compiled Edge TPU model.
+        prefix (str, optional): Logging prefix. Defaults to "".
+
+    Notes:
+        Requires the Edge TPU compiler to be installed. The function compiles the TFLite model
+        for optimal performance on Google's Edge TPU hardware accelerator.
+    """
+    import subprocess
+
+    cmd = (
+        "edgetpu_compiler "
+        f'--out_dir "{output_dir}" '
+        "--show_operations "
+        "--search_delegate "
+        "--delegate_search_step 30 "
+        "--timeout_sec 180 "
+        f'"{tflite_file}"'
+    )
+    LOGGER.info(f"{prefix} running '{cmd}'")
+    subprocess.run(cmd, shell=True)
+
+
+def pb2tfjs(pb_file: str, output_dir: str, half: bool = False, int8: bool = False, prefix: str = ""):
+    """Convert a TensorFlow GraphDef (.pb) model to TensorFlow.js format.
+
+    Args:
+        pb_file (str): Path to the input TensorFlow GraphDef (.pb) model file.
+        output_dir (str): Output directory path for the converted TensorFlow.js model.
+        half (bool, optional): Enable FP16 quantization. Defaults to False.
+        int8 (bool, optional): Enable INT8 quantization. Defaults to False.
+        prefix (str, optional): Logging prefix. Defaults to "".
+
+    Notes:
+        Requires tensorflowjs package. Uses tensorflowjs_converter command-line tool for conversion.
+        Handles spaces in file paths and warns if output directory contains spaces.
+    """
+    import subprocess
+
+    import tensorflow as tf
+    import tensorflowjs as tfjs
+
+    LOGGER.info(f"\n{prefix} starting export with tensorflowjs {tfjs.__version__}...")
+
+    gd = tf.Graph().as_graph_def()  # TF GraphDef
+    with open(pb_file, "rb") as file:
+        gd.ParseFromString(file.read())
+    outputs = ",".join(gd_outputs(gd))
+    LOGGER.info(f"\n{prefix} output node names: {outputs}")
+
+    quantization = "--quantize_float16" if half else "--quantize_uint8" if int8 else ""
+    with spaces_in_path(pb_file) as fpb_, spaces_in_path(output_dir) as f_:  # exporter cannot handle spaces in paths
+        cmd = (
+            "tensorflowjs_converter "
+            f'--input_format=tf_frozen_model {quantization} --output_node_names={outputs} "{fpb_}" "{f_}"'
+        )
+        LOGGER.info(f"{prefix} running '{cmd}'")
+        subprocess.run(cmd, shell=True)
+
+    if " " in output_dir:
+        LOGGER.warning(f"{prefix} your model may not work correctly with spaces in path '{output_dir}'.")
+
+
+def gd_outputs(gd):
+    """Return TensorFlow GraphDef model output node names."""
+    name_list, input_list = [], []
+    for node in gd.node:  # tensorflow.core.framework.node_def_pb2.NodeDef
+        name_list.append(node.name)
+        input_list.extend(node.input)
+    return sorted(f"{x}:0" for x in list(set(name_list) - set(input_list)) if not x.startswith("NoOp"))

ultralytics/utils/files.py

@@ -13,11 +13,10 @@ from pathlib import Path
 
 
 class WorkingDirectory(contextlib.ContextDecorator):
-    """
-    A context manager and decorator for temporarily changing the working directory.
+    """A context manager and decorator for temporarily changing the working directory.
 
-    This class allows for the temporary change of the working directory using a context manager or decorator.
-    It ensures that the original working directory is restored after the context or decorated function completes.
+    This class allows for the temporary change of the working directory using a context manager or decorator. It ensures
+    that the original working directory is restored after the context or decorated function completes.
 
     Attributes:
         dir (Path | str): The new directory to switch to.
@@ -29,15 +28,15 @@ class WorkingDirectory(contextlib.ContextDecorator):
 
     Examples:
         Using as a context manager:
-        >>> with WorkingDirectory('/path/to/new/dir'):
-        >>> # Perform operations in the new directory
-        >>>     pass
+        >>> with WorkingDirectory("/path/to/new/dir"):
+        ...     # Perform operations in the new directory
+        ...     pass
 
         Using as a decorator:
-        >>> @WorkingDirectory('/path/to/new/dir')
-        >>> def some_function():
-        >>> # Perform operations in the new directory
-        >>>     pass
+        >>> @WorkingDirectory("/path/to/new/dir")
+        ... def some_function():
+        ...     # Perform operations in the new directory
+        ...     pass
     """
 
     def __init__(self, new_dir: str | Path):
@@ -49,15 +48,14 @@ class WorkingDirectory(contextlib.ContextDecorator):
         """Change the current working directory to the specified directory upon entering the context."""
         os.chdir(self.dir)
 
-    def __exit__(self, exc_type, exc_val, exc_tb):  # noqa
+    def __exit__(self, exc_type, exc_val, exc_tb):
         """Restore the original working directory when exiting the context."""
         os.chdir(self.cwd)
 
 
 @contextmanager
 def spaces_in_path(path: str | Path):
-    """
-    Context manager to handle paths with spaces in their names.
+    """Context manager to handle paths with spaces in their names.
 
     If a path contains spaces, it replaces them with underscores, copies the file/directory to the new path, executes
     the context code block, then copies the file/directory back to its original location.
@@ -66,13 +64,12 @@ def spaces_in_path(path: str | Path):
         path (str | Path): The original path that may contain spaces.
 
     Yields:
-        (Path | str): Temporary path with spaces replaced by underscores if spaces were present, otherwise the
-            original path.
+        (Path | str): Temporary path with any spaces replaced by underscores.
 
     Examples:
-        >>> with spaces_in_path('/path/with spaces') as new_path:
-        >>> # Your code here
-        >>>     pass
+        >>> with spaces_in_path("/path/with spaces") as new_path:
+        ...     # Your code here
+        ...     pass
     """
     # If path has spaces, replace them with underscores
     if " " in str(path):
@@ -107,12 +104,11 @@ def spaces_in_path(path: str | Path):
 
 
 def increment_path(path: str | Path, exist_ok: bool = False, sep: str = "", mkdir: bool = False) -> Path:
-    """
-    Increment a file or directory path, i.e., runs/exp --> runs/exp{sep}2, runs/exp{sep}3, ... etc.
+    """Increment a file or directory path, i.e., runs/exp --> runs/exp{sep}2, runs/exp{sep}3, ... etc.
 
-    If the path exists and `exist_ok` is not True, the path will be incremented by appending a number and `sep` to
-    the end of the path. If the path is a file, the file extension will be preserved. If the path is a directory, the
-    number will be appended directly to the end of the path.
+    If the path exists and `exist_ok` is not True, the path will be incremented by appending a number and `sep` to the
+    end of the path. If the path is a file, the file extension will be preserved. If the path is a directory, the number
+    will be appended directly to the end of the path.
 
     Args:
         path (str | Path): Path to increment.
@@ -185,8 +181,7 @@ def get_latest_run(search_dir: str = ".") -> str:
 
 
 def update_models(model_names: tuple = ("yolo11n.pt",), source_dir: Path = Path("."), update_names: bool = False):
-    """
-    Update and re-save specified YOLO models in an 'updated_models' subdirectory.
+    """Update and re-save specified YOLO models in an 'updated_models' subdirectory.
 
     Args:
         model_names (tuple, optional): Model filenames to update.
@@ -201,13 +196,14 @@ def update_models(model_names: tuple = ("yolo11n.pt",), source_dir: Path = Path(
     """
     from ultralytics import YOLO
     from ultralytics.nn.autobackend import default_class_names
+    from ultralytics.utils import LOGGER
 
     target_dir = source_dir / "updated_models"
     target_dir.mkdir(parents=True, exist_ok=True)  # Ensure target directory exists
 
     for model_name in model_names:
         model_path = source_dir / model_name
-        print(f"Loading model from {model_path}")
+        LOGGER.info(f"Loading model from {model_path}")
 
         # Load model
         model = YOLO(model_path)
@@ -219,5 +215,5 @@ def update_models(model_names: tuple = ("yolo11n.pt",), source_dir: Path = Path(
         save_path = target_dir / model_name
 
         # Save model using model.save()
-        print(f"Re-saving {model_name} model to {save_path}")
+        LOGGER.info(f"Re-saving {model_name} model to {save_path}")
         model.save(save_path)