dgenerate-ultralytics-headless 8.3.221__py3-none-any.whl → 8.3.223__py3-none-any.whl

ultralytics/utils/export/imx.py

@@ -6,6 +6,7 @@ import subprocess
 import types
 from pathlib import Path
 
+import numpy as np
 import torch
 
 from ultralytics.nn.modules import Detect, Pose
@@ -13,6 +14,32 @@ from ultralytics.utils import LOGGER
 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},
+    },
+    "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},
+    },
+}
+
 
 class FXModel(torch.nn.Module):
     """
@@ -200,30 +227,13 @@ def torch2imx(
     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 +242,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,13 +265,14 @@ 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)

ultralytics/utils/export/tensorflow.py

@@ -0,0 +1,221 @@
+# 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.
+
+    Note:
+        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()
+            np.save(str(tmp_file), images)  # BHWC
+            np_data = [["images", tmp_file, [[[[0, 0, 0]]]], [[[[255, 255, 255]]]]]]
+
+    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(tf_keras): 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 "".
+
+    Note:
+        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 "".
+
+    Note:
+        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 "".
+
+    Note:
+        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 can not handle spaces in path
+        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/metrics.py

@@ -663,7 +663,7 @@ def plot_pr_curve(
         for i, y in enumerate(py.T):
             ax.plot(px, y, linewidth=1, label=f"{names[i]} {ap[i, 0]:.3f}")  # plot(recall, precision)
     else:
-        ax.plot(px, py, linewidth=1, color="grey")  # plot(recall, precision)
+        ax.plot(px, py, linewidth=1, color="gray")  # plot(recall, precision)
 
     ax.plot(px, py.mean(1), linewidth=3, color="blue", label=f"all classes {ap[:, 0].mean():.3f} mAP@0.5")
     ax.set_xlabel("Recall")
@@ -708,7 +708,7 @@ def plot_mc_curve(
         for i, y in enumerate(py):
             ax.plot(px, y, linewidth=1, label=f"{names[i]}")  # plot(confidence, metric)
     else:
-        ax.plot(px, py.T, linewidth=1, color="grey")  # plot(confidence, metric)
+        ax.plot(px, py.T, linewidth=1, color="gray")  # plot(confidence, metric)
 
     y = smooth(py.mean(0), 0.1)
     ax.plot(px, y, linewidth=3, color="blue", label=f"all classes {y.max():.2f} at {px[y.argmax()]:.3f}")

ultralytics/utils/nms.py

@@ -231,9 +231,11 @@ class TorchNMS:
             upper_mask = row_idx < col_idx
             ious = ious * upper_mask
             # Zeroing these scores ensures the additional indices would not affect the final results
-            scores[~((ious >= iou_threshold).sum(0) <= 0)] = 0
+            scores_ = scores[sorted_idx]
+            scores_[~((ious >= iou_threshold).sum(0) <= 0)] = 0
+            scores[sorted_idx] = scores_  # update original tensor for NMSModel
             # NOTE: return indices with fixed length to avoid TFLite reshape error
-            pick = torch.topk(scores, scores.shape[0]).indices
+            pick = torch.topk(scores_, scores_.shape[0]).indices
         return sorted_idx[pick]
 
     @staticmethod

ultralytics/utils/plotting.py

@@ -722,7 +722,7 @@ def plot_images(
         convert tensor inputs to numpy arrays for processing.
 
         Channel Support:
-        - 1 channel: Greyscale
+        - 1 channel: Grayscale
         - 2 channels: Third channel added as zeros
         - 3 channels: Used as-is (standard RGB)
         - 4+ channels: Cropped to first 3 channels