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

ultralytics/utils/export/imx.py

@@ -21,27 +21,39 @@ from ultralytics.utils.torch_utils import copy_attr
 MCT_CONFIG = {
     "YOLO11": {
         "detect": {
-            "layer_names": ["sub", "mul_2", "add_14", "cat_21"],
+            "layer_names": ["sub", "mul_2", "add_14", "cat_19"],
             "weights_memory": 2585350.2439,
-            "n_layers": 238,
+            "n_layers": {238, 239},
         },
         "pose": {
-            "layer_names": ["sub", "mul_2", "add_14", "cat_22", "cat_23", "mul_4", "add_15"],
+            "layer_names": ["sub", "mul_2", "add_14", "cat_21", "cat_22", "mul_4", "add_15"],
             "weights_memory": 2437771.67,
-            "n_layers": 257,
+            "n_layers": {257, 258},
+        },
+        "classify": {"layer_names": [], "weights_memory": np.inf, "n_layers": {112}},
+        "segment": {
+            "layer_names": ["sub", "mul_2", "add_14", "cat_21"],
+            "weights_memory": 2466604.8,
+            "n_layers": {265, 266},
         },
-        "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},
+        "detect": {
+            "layer_names": ["sub", "mul", "add_6", "cat_15"],
+            "weights_memory": 2550540.8,
+            "n_layers": {168, 169},
+        },
         "pose": {
-            "layer_names": ["add_7", "mul_2", "cat_19", "mul", "sub", "add_6", "cat_18"],
+            "layer_names": ["add_7", "mul_2", "cat_17", "mul", "sub", "add_6", "cat_18"],
             "weights_memory": 2482451.85,
-            "n_layers": 187,
+            "n_layers": {187, 188},
+        },
+        "classify": {"layer_names": [], "weights_memory": np.inf, "n_layers": {73}},
+        "segment": {
+            "layer_names": ["sub", "mul", "add_6", "cat_17"],
+            "weights_memory": 2580060.0,
+            "n_layers": {195, 196},
         },
-        "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},
     },
 }
 
@@ -104,20 +116,26 @@ class FXModel(torch.nn.Module):
         return x
 
 
-def _inference(self, x: list[torch.Tensor]) -> tuple[torch.Tensor]:
+def _inference(self, x: dict[str, torch.Tensor]) -> tuple[torch.Tensor]:
     """Decode boxes and cls scores for imx object detection."""
-    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)
-    dbox = self.decode_bboxes(self.dfl(box), self.anchors.unsqueeze(0)) * self.strides
-    return dbox.transpose(1, 2), cls.sigmoid().permute(0, 2, 1)
+    dbox = self.decode_bboxes(self.dfl(x["boxes"]), self.anchors.unsqueeze(0)) * self.strides
+    return dbox.transpose(1, 2), x["scores"].sigmoid().permute(0, 2, 1)
 
 
 def pose_forward(self, x: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """Forward pass for imx pose estimation, including keypoint decoding."""
     bs = x[0].shape[0]  # batch size
-    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)
+    nk_out = getattr(self, "nk_output", self.nk)
+    kpt = torch.cat([self.cv4[i](x[i]).view(bs, nk_out, -1) for i in range(self.nl)], -1)
+
+    # If using Pose26 with 5 dims, convert to 3 dims for export
+    if hasattr(self, "nk_output") and self.nk_output != self.nk:
+        spatial = kpt.shape[-1]
+        kpt = kpt.view(bs, self.kpt_shape[0], self.kpt_shape[1] + 2, spatial)
+        kpt = kpt[:, :, :-2, :]  # Remove sigma_x, sigma_y
+        kpt = kpt.view(bs, self.nk, spatial)
     x = Detect.forward(self, x)
-    pred_kpt = self.kpts_decode(bs, kpt)
+    pred_kpt = self.kpts_decode(kpt)
     return *x, pred_kpt.permute(0, 2, 1)
 
 
@@ -219,7 +237,7 @@ def torch2imx(
     Examples:
         >>> from ultralytics import YOLO
         >>> model = YOLO("yolo11n.pt")
-        >>> path, _ = export_imx(model, "model.imx", conf=0.25, iou=0.45, max_det=300)
+        >>> path, _ = export_imx(model, "model.imx", conf=0.25, iou=0.7, max_det=300)
 
     Notes:
         - Requires model_compression_toolkit, onnx, edgemdt_tpc, and edge-mdt-cl packages
@@ -245,7 +263,7 @@ def torch2imx(
     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())) != mct_config["n_layers"]:
+    if len(list(model.modules())) not in mct_config["n_layers"]:
         raise ValueError("IMX export only supported for YOLOv8n and YOLO11n models.")
 
     for layer_name in mct_config["layer_names"]:

ultralytics/utils/export/tensorflow.py

@@ -2,12 +2,13 @@
 
 from __future__ import annotations
 
+from functools import partial
 from pathlib import Path
 
 import numpy as np
 import torch
 
-from ultralytics.nn.modules import Detect, Pose
+from ultralytics.nn.modules import Detect, Pose, Pose26
 from ultralytics.utils import LOGGER
 from ultralytics.utils.downloads import attempt_download_asset
 from ultralytics.utils.files import spaces_in_path
@@ -15,43 +16,42 @@ 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."""
+    """A wrapper for TensorFlow export compatibility (TF-specific handling is now in head modules)."""
     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)
+        m._get_decode_boxes = types.MethodType(_tf_decode_boxes, m)
+        if isinstance(m, Pose):
+            m.kpts_decode = types.MethodType(partial(_tf_kpts_decode, is_pose26=type(m) is Pose26), 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))
+def _tf_decode_boxes(self, x: dict[str, torch.Tensor]) -> torch.Tensor:
+    """Decode bounding boxes for TensorFlow export."""
+    shape = x["feats"][0].shape  # BCHW
+    boxes = x["boxes"]
+    if self.format != "imx" and (self.dynamic or self.shape != shape):
+        self.anchors, self.strides = (a.transpose(0, 1) for a in make_anchors(x["feats"], 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)
+    grid_h, grid_w = shape[2:4]
+    grid_size = torch.tensor([grid_w, grid_h, grid_w, grid_h], device=boxes.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)
+    dbox = self.decode_bboxes(self.dfl(boxes) * norm, self.anchors.unsqueeze(0) * norm[:, :2])
+    return dbox
 
 
-def tf_kpts_decode(self, bs: int, kpts: torch.Tensor) -> torch.Tensor:
-    """Decode keypoints for tf pose estimation."""
+def _tf_kpts_decode(self, kpts: torch.Tensor, is_pose26: bool = False) -> torch.Tensor:
+    """Decode keypoints for TensorFlow export."""
     ndim = self.kpt_shape[1]
-    # required for TFLite export to avoid 'PLACEHOLDER_FOR_GREATER_OP_CODES' bug
+    bs = kpts.shape[0]
     # 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_h, grid_w = self.shape[2:4]
     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
+    a = ((y[:, :, :2] + self.anchors) if is_pose26 else (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)

ultralytics/utils/files.py

@@ -180,7 +180,7 @@ def get_latest_run(search_dir: str = ".") -> str:
     return max(last_list, key=os.path.getctime) if last_list else ""
 
 
-def update_models(model_names: tuple = ("yolo11n.pt",), source_dir: Path = Path("."), update_names: bool = False):
+def update_models(model_names: tuple = ("yolo26n.pt",), source_dir: Path = Path("."), update_names: bool = False):
     """Update and re-save specified YOLO models in an 'updated_models' subdirectory.
 
     Args:
@@ -191,7 +191,7 @@ def update_models(model_names: tuple = ("yolo11n.pt",), source_dir: Path = Path(
     Examples:
         Update specified YOLO models and save them in 'updated_models' subdirectory:
         >>> from ultralytics.utils.files import update_models
-        >>> model_names = ("yolo11n.pt", "yolov8s.pt")
+        >>> model_names = ("yolo26n.pt", "yolo11s.pt")
         >>> update_models(model_names, source_dir=Path("/models"), update_names=True)
     """
     from ultralytics import YOLO