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

{dgenerate_ultralytics_headless-8.3.222.dist-info → dgenerate_ultralytics_headless-8.3.223.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dgenerate-ultralytics-headless
-Version: 8.3.222
+Version: 8.3.223
 Summary: Automatically built Ultralytics package with python-opencv-headless dependency instead of python-opencv
 Author-email: Glenn Jocher <glenn.jocher@ultralytics.com>, Jing Qiu <jing.qiu@ultralytics.com>
 Maintainer-email: Ultralytics <hello@ultralytics.com>
@@ -44,7 +44,7 @@ Requires-Dist: torch!=2.4.0,>=1.8.0; sys_platform == "win32"
 Requires-Dist: torchvision>=0.9.0
 Requires-Dist: psutil
 Requires-Dist: polars
-Requires-Dist: ultralytics-thop>=2.0.0
+Requires-Dist: ultralytics-thop>=2.0.18
 Provides-Extra: dev
 Requires-Dist: ipython; extra == "dev"
 Requires-Dist: pytest; extra == "dev"

{dgenerate_ultralytics_headless-8.3.222.dist-info → dgenerate_ultralytics_headless-8.3.223.dist-info}/RECORD

@@ -1,4 +1,4 @@
-dgenerate_ultralytics_headless-8.3.222.dist-info/licenses/LICENSE,sha256=DZak_2itbUtvHzD3E7GNUYSRK6jdOJ-GqncQ2weavLA,34523
+dgenerate_ultralytics_headless-8.3.223.dist-info/licenses/LICENSE,sha256=DZak_2itbUtvHzD3E7GNUYSRK6jdOJ-GqncQ2weavLA,34523
 tests/__init__.py,sha256=bCox_hLdGRFYGLb2kd722VdNP2zEXNYNuLLYtqZSrbw,804
 tests/conftest.py,sha256=oaqn_-8LH7R4YQAKKOiK3iuAfjmmLQ9-pL_IPj6xq-U,2333
 tests/test_cli.py,sha256=zygPlaksok7Nwugp3aIudDSkOlzISvmDWfKNmpY3mSA,5844
@@ -8,7 +8,7 @@ tests/test_exports.py,sha256=OMLio2uUhyqo8D8qB5xUwmk7Po2rMeAACRc8WYoxbj4,13147
 tests/test_integrations.py,sha256=6QgSh9n0J04RdUYz08VeVOnKmf4S5MDEQ0chzS7jo_c,6220
 tests/test_python.py,sha256=OChceQcDDAy07yACnmOoGfimRo_4YdyiMwukGEgozXA,27735
 tests/test_solutions.py,sha256=j_PZZ5tMR1Y5ararY-OTXZr1hYJ7vEVr8H3w4O1tbQs,14153
-ultralytics/__init__.py,sha256=pcHeLhMCBEuP-sM307qncLjoxjr7uW1KqDHyi7RwD60,1302
+ultralytics/__init__.py,sha256=IFuXT77f7jmVOvOHnLjLEIrgQ-RfhI6Rq7ykdDC42GI,1302
 ultralytics/py.typed,sha256=la67KBlbjXN-_-DfGNcdOcjYumVpKG_Tkw-8n5dnGB4,8
 ultralytics/assets/bus.jpg,sha256=wCAZxJecGR63Od3ZRERe9Aja1Weayrb9Ug751DS_vGM,137419
 ultralytics/assets/zidane.jpg,sha256=Ftc4aeMmen1O0A3o6GCDO9FlfBslLpTAw0gnetx7bts,50427
@@ -19,7 +19,7 @@ ultralytics/cfg/datasets/DOTAv1.5.yaml,sha256=VZ_KKFX0H2YvlFVJ8JHcLWYBZ2xiQ6Z-RO
 ultralytics/cfg/datasets/DOTAv1.yaml,sha256=JrDuYcQ0JU9lJlCA-dCkMNko_jaj6MAVGHjsfjeZ_u0,1181
 ultralytics/cfg/datasets/GlobalWheat2020.yaml,sha256=dnr_loeYSE6Eo_f7V1yubILsMRBMRm1ozyC5r7uT-iY,2144
 ultralytics/cfg/datasets/HomeObjects-3K.yaml,sha256=xEtSqEad-rtfGuIrERjjhdISggmPlvaX-315ZzKz50I,934
-ultralytics/cfg/datasets/ImageNet.yaml,sha256=GvDWypLVG_H3H67Ai8IC1pvK6fwcTtF5FRhzO1OXXDU,42530
+ultralytics/cfg/datasets/ImageNet.yaml,sha256=N9NHhIgnlNIBqZZbzQZAW3aCnz6RSXQABnopaDs5BmE,42529
 ultralytics/cfg/datasets/Objects365.yaml,sha256=8Bl-NAm0mlMW8EfMsz39JZo-HCvmp0ejJXaMeoHTpqw,9649
 ultralytics/cfg/datasets/SKU-110K.yaml,sha256=xvRkq3SdDOwBA91U85bln7HTXkod5MvFX6pt1PxTjJE,2609
 ultralytics/cfg/datasets/VOC.yaml,sha256=84BaL-iwG03M_W9hNzjgEQi918dZgSHbCgf9DShjwLA,3747
@@ -42,9 +42,9 @@ ultralytics/cfg/datasets/dog-pose.yaml,sha256=BI-2S3_cSVyV2Gfzbs_3GzvivRlikT0ANj
 ultralytics/cfg/datasets/dota8-multispectral.yaml,sha256=2lMBi1Q3_pc0auK00yX80oF7oUMo0bUlwjkOrp33hvs,1216
 ultralytics/cfg/datasets/dota8.yaml,sha256=5n4h_4zdrtUSkmH5DHJ-JLPvfiATcieIkgP3NeOP5nI,1060
 ultralytics/cfg/datasets/hand-keypoints.yaml,sha256=NglEDsfNRe0DaYnwy7n6hYUxEAjV-V2NZBUbj1qJaag,1365
-ultralytics/cfg/datasets/lvis.yaml,sha256=lMvPfuiDv_o2qLxAWoh9WMrvjKJ5moLrcx1gr3RG_pM,29680
+ultralytics/cfg/datasets/lvis.yaml,sha256=RescdwAJ8EU1o7Sm0YlxYsGbQFNU1p-LFbFKYEt5MhE,29596
 ultralytics/cfg/datasets/medical-pills.yaml,sha256=RK7iQFpDDkUS6EsEGqlbFjoohi3cgSsUIbsk7UItyds,792
-ultralytics/cfg/datasets/open-images-v7.yaml,sha256=wK9v3OAGdHORkFdqoBi0hS0fa1b74LLroAzUSWjxEqw,12119
+ultralytics/cfg/datasets/open-images-v7.yaml,sha256=2fVFmb8UEYH-LkX0z5GlYp__U0_GDqVgVqzmnfFerm8,12116
 ultralytics/cfg/datasets/package-seg.yaml,sha256=V4uyTDWWzgft24y9HJWuELKuZ5AndAHXbanxMI6T8GU,849
 ultralytics/cfg/datasets/signature.yaml,sha256=gBvU3715gVxVAafI_yaYczGX3kfEfA4BttbiMkgOXNk,774
 ultralytics/cfg/datasets/tiger-pose.yaml,sha256=bJ7nBTDQwXRHtlg3xmo4C2bOpPn_r4l8-DezSWMYNcU,1196
@@ -122,7 +122,7 @@ ultralytics/data/scripts/get_coco.sh,sha256=UuJpJeo3qQpTHVINeOpmP0NYmg8PhEFE3A8J
 ultralytics/data/scripts/get_coco128.sh,sha256=qmRQl_hOKrsdHrTrnyQuFIH01oDz3lfaz138OgGfLt8,650
 ultralytics/data/scripts/get_imagenet.sh,sha256=hr42H16bM47iT27rgS7MpEo-GeOZAYUQXgr0B2cwn48,1705
 ultralytics/engine/__init__.py,sha256=lm6MckFYCPTbqIoX7w0s_daxdjNeBeKW6DXppv1-QUM,70
-ultralytics/engine/exporter.py,sha256=zFk58rZvZvaxK9dIwHOLjnqcoyRnQn0l5t2RbNIGsIE,73417
+ultralytics/engine/exporter.py,sha256=89hggNbcH7zFAG8QJmShoHFZMvn0SpHF_yTEJ4CMbsc,69852
 ultralytics/engine/model.py,sha256=d7yGl8ybd7v8W4Q-ueSDAVfumDhsx0QCp4mx8OKf0Z8,53448
 ultralytics/engine/predictor.py,sha256=ZQrx1Bz4X8aTgGjrOSdRSP7SCtQ05uqz6IitEan_Gyk,22813
 ultralytics/engine/results.py,sha256=oHQdV_eIMvAU2qLCV7wG7iLifdfaLEgP80lDXB5ghkg,71490
@@ -148,7 +148,7 @@ ultralytics/models/rtdetr/__init__.py,sha256=F4NEQqtcVKFxj97Dh7rkn2Vu3JG4Ea_nxqr
 ultralytics/models/rtdetr/model.py,sha256=Pq9QDgaZetDnjxdYSoomj2s6vOGSdpsqVfyN5j0GUmc,2292
 ultralytics/models/rtdetr/predict.py,sha256=43-gGCHEH7UQQ6H1oXdlDlrM39esnp-YEhqCvZOwtOM,4279
 ultralytics/models/rtdetr/train.py,sha256=SNntxGHXatbNqn1yna5_dDQiR_ciDK6o_4S7JIHU7EY,3765
-ultralytics/models/rtdetr/val.py,sha256=l26CzpcYHYC0sQ--rKUFBCYl73nsgAGOj1U3xScNzFs,8918
+ultralytics/models/rtdetr/val.py,sha256=UXaoNiy81zdkv6d79x1oGyR8T7dwuV5Y4m0Gpe-LQts,8976
 ultralytics/models/sam/__init__.py,sha256=p1BKLawQFvVxmdk7LomFVWX-67Kc-AP4PJBNPfU_Nuc,359
 ultralytics/models/sam/amg.py,sha256=nFq4EwHf65W2N5Ipo4W69nGRhCbJEh_boYQ8SIPWBZ0,11816
 ultralytics/models/sam/build.py,sha256=uKCgHpcYgV26OFuMq5RaGR8aXYoEtNoituT06bmnW44,12790
@@ -197,14 +197,14 @@ ultralytics/models/yolo/yoloe/train.py,sha256=qefvNNXDTOK1tO3va0kNHr8lE5QJkOlV8G
 ultralytics/models/yolo/yoloe/train_seg.py,sha256=aCV7M8oQOvODFnU4piZdJh3tIrBJYAzZfRVRx1vRgxo,4956
 ultralytics/models/yolo/yoloe/val.py,sha256=5Gd9EoFH0FmKKvWXBl4J7gBe9DVxIczN-s3ceHwdUDo,9458
 ultralytics/nn/__init__.py,sha256=538LZPUKKvc3JCMgiQ4VLGqRN2ZAaVLFcQbeNNHFkEA,545
-ultralytics/nn/autobackend.py,sha256=918iNweM3fTuRIbxHbXC1wspOj9rlGkuwalQ61uYLbk,42694
+ultralytics/nn/autobackend.py,sha256=gw8REfburF36l9Hyh11eYzy7UnMvuX1Dm3cjsJBA1TM,42702
 ultralytics/nn/tasks.py,sha256=vRr6HTucM7Eg3kxzhYtyjgEAdacZ7gIDU3yPbMnyYMM,70834
 ultralytics/nn/text_model.py,sha256=pHqnKe8UueR1MuwJcIE_IvrnYIlt68QL796xjcRJs2A,15275
 ultralytics/nn/modules/__init__.py,sha256=5Sg_28MDfKwdu14Ty_WCaiIXZyjBSQ-xCNCwnoz_w-w,3198
 ultralytics/nn/modules/activation.py,sha256=75JcIMH2Cu9GTC2Uf55r_5YLpxcrXQDaVoeGQ0hlUAU,2233
 ultralytics/nn/modules/block.py,sha256=eQ8DegyvBG9k-O_QgSZe5XGmpravqwlnSCCBW6bHRXo,70622
 ultralytics/nn/modules/conv.py,sha256=MISNAK8NzAZhNUusVKWvTHQ8IsofwM-5X0gChCagsaY,21457
-ultralytics/nn/modules/head.py,sha256=HBSoHOcd9hikpBLF9BdgIItYJDNWIO8NiNYMcMw6ThM,53512
+ultralytics/nn/modules/head.py,sha256=XBOLfpxgApIhNmdgnWoECep0wKhrw8LWtmd1TrWNBak,52076
 ultralytics/nn/modules/transformer.py,sha256=9aq0Yo9V3C4y_McSje4qE1d_PTWDctTsrb98MyXxigc,31470
 ultralytics/nn/modules/utils.py,sha256=9kLeEtvEBFLugz53TkdI4mifD-39a-upjPD-wrE8opU,6092
 ultralytics/solutions/__init__.py,sha256=Jj7OcRiYjHH-e104H4xTgjjR5W6aPB4mBRndbaSPmgU,1209
@@ -253,7 +253,7 @@ ultralytics/utils/instance.py,sha256=_b_jMTECWJGzncCiTg7FtTDSSeXGnbiAhaJhIsqbn9k
 ultralytics/utils/logger.py,sha256=hK1APBBHmlLAm0zbAFY7gf7Iaejy0PdwLWnnpboboGg,15129
 ultralytics/utils/loss.py,sha256=wJ0F2DpRTI9-e9adxIm2io0zcXRa0RTWFTOc7WmS1-A,39827
 ultralytics/utils/metrics.py,sha256=EWwkVWNmN_9rIsR1UOTLz3PiXOzflUE0iWFibydvXgM,68882
-ultralytics/utils/nms.py,sha256=AVOmPuUTEJqmq2J6rvjq-nHNxYIyabgzHdc41siyA0w,14161
+ultralytics/utils/nms.py,sha256=SnZF0VRzY933YzI92NLzmLwuVzu56UNZ7sFT0FryCaw,14285
 ultralytics/utils/ops.py,sha256=yb0jlahjxqUT_xb3y9wz0kXn0rx2AryUgWdtLat3yWY,27010
 ultralytics/utils/patches.py,sha256=0-2G4jXCIPnMonlft-cPcjfFcOXQS6ODwUDNUwanfg4,6541
 ultralytics/utils/plotting.py,sha256=l5G4MT2pB_LLMFqSgFbKb7ip5VMrnpi3i5QmZWytRRU,48369
@@ -274,10 +274,12 @@ ultralytics/utils/callbacks/platform.py,sha256=a7T_8htoBB0uX1WIc392UJnhDjxkRyQMv
 ultralytics/utils/callbacks/raytune.py,sha256=S6Bq16oQDQ8BQgnZzA0zJHGN_BBr8iAM_WtGoLiEcwg,1283
 ultralytics/utils/callbacks/tensorboard.py,sha256=_4nfGK1dDLn6ijpvphBDhc-AS8qhS3jjY2CAWB7SNF0,5283
 ultralytics/utils/callbacks/wb.py,sha256=ngQO8EJ1kxJDF1YajScVtzBbm26jGuejA0uWeOyvf5A,7685
-ultralytics/utils/export/__init__.py,sha256=uyRhb-0Z5FVf7vSz2Yba1m7g5m2U_ftAv4ThlmMsqZ8,10015
+ultralytics/utils/export/__init__.py,sha256=Cfh-PwVfTF_lwPp-Ss4wiX4z8Sm1XRPklsqdFfmTZ30,333
+ultralytics/utils/export/engine.py,sha256=V8ERERlpufTRm6k_7KOy9dUupAICC28W9TPO_7dkEJY,9979
 ultralytics/utils/export/imx.py,sha256=DH0rVe-gris7qA7bGT-WoOJHqWxLBAmei1JXmK-W7vM,11660
-dgenerate_ultralytics_headless-8.3.222.dist-info/METADATA,sha256=ebl0D5CHZMeyOJjM5oKV2O9pGBocaEfo8x0MzGFQXEs,38763
-dgenerate_ultralytics_headless-8.3.222.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-dgenerate_ultralytics_headless-8.3.222.dist-info/entry_points.txt,sha256=YM_wiKyTe9yRrsEfqvYolNO5ngwfoL4-NwgKzc8_7sI,93
-dgenerate_ultralytics_headless-8.3.222.dist-info/top_level.txt,sha256=XP49TwiMw4QGsvTLSYiJhz1xF_k7ev5mQ8jJXaXi45Q,12
-dgenerate_ultralytics_headless-8.3.222.dist-info/RECORD,,
+ultralytics/utils/export/tensorflow.py,sha256=CxraBn-5pIDSd_-0-0vQGMz8lv75vjSl6N7DYgVS3SU,9382
+dgenerate_ultralytics_headless-8.3.223.dist-info/METADATA,sha256=P2UOqD5tY6Tx3cfHaU_B5oXQPOsycKUo6lThBHn7G0s,38764
+dgenerate_ultralytics_headless-8.3.223.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+dgenerate_ultralytics_headless-8.3.223.dist-info/entry_points.txt,sha256=YM_wiKyTe9yRrsEfqvYolNO5ngwfoL4-NwgKzc8_7sI,93
+dgenerate_ultralytics_headless-8.3.223.dist-info/top_level.txt,sha256=XP49TwiMw4QGsvTLSYiJhz1xF_k7ev5mQ8jJXaXi45Q,12
+dgenerate_ultralytics_headless-8.3.223.dist-info/RECORD,,

ultralytics/__init__.py

@@ -1,6 +1,6 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
-__version__ = "8.3.222"
+__version__ = "8.3.223"
 
 import importlib
 import os

ultralytics/cfg/datasets/ImageNet.yaml

@@ -342,7 +342,7 @@ names:
   322: ringlet
   323: monarch butterfly
   324: small white
-  325: sulphur butterfly
+  325: sulfur butterfly
   326: gossamer-winged butterfly
   327: starfish
   328: sea urchin

ultralytics/cfg/datasets/lvis.yaml

@@ -35,7 +35,7 @@ names:
   17: armband
   18: armchair
   19: armoire
-  20: armor/armour
+  20: armor
   21: artichoke
   22: trash can/garbage can/wastebin/dustbin/trash barrel/trash bin
   23: ashtray
@@ -245,7 +245,7 @@ names:
   227: CD player
   228: celery
   229: cellular telephone/cellular phone/cellphone/mobile phone/smart phone
-  230: chain mail/ring mail/chain armor/chain armour/ring armor/ring armour
+  230: chain mail/ring mail/chain armor/ring armor
   231: chair
   232: chaise longue/chaise/daybed
   233: chalice
@@ -305,7 +305,7 @@ names:
   287: coin
   288: colander/cullender
   289: coleslaw/slaw
-  290: coloring material/colouring material
+  290: coloring material
   291: combination lock
   292: pacifier/teething ring
   293: comic book
@@ -401,7 +401,7 @@ names:
   383: domestic ass/donkey
   384: doorknob/doorhandle
   385: doormat/welcome mat
-  386: doughnut/donut
+  386: donut
   387: dove
   388: dragonfly
   389: drawer
@@ -1072,7 +1072,7 @@ names:
   1054: tag
   1055: taillight/rear light
   1056: tambourine
-  1057: army tank/armored combat vehicle/armoured combat vehicle
+  1057: army tank/armored combat vehicle
   1058: tank/tank storage vessel/storage tank
   1059: tank top/tank top clothing
   1060: tape/tape sticky cloth or paper

ultralytics/cfg/datasets/open-images-v7.yaml

@@ -182,7 +182,7 @@ names:
   163: Dolphin
   164: Door
   165: Door handle
-  166: Doughnut
+  166: Donut
   167: Dragonfly
   168: Drawer
   169: Dress

ultralytics/engine/exporter.py

@@ -107,9 +107,17 @@ from ultralytics.utils.checks import (
     is_intel,
     is_sudo_available,
 )
-from ultralytics.utils.downloads import attempt_download_asset, get_github_assets, safe_download
-from ultralytics.utils.export import onnx2engine, torch2imx, torch2onnx
-from ultralytics.utils.files import file_size, spaces_in_path
+from ultralytics.utils.downloads import get_github_assets, safe_download
+from ultralytics.utils.export import (
+    keras2pb,
+    onnx2engine,
+    onnx2saved_model,
+    pb2tfjs,
+    tflite2edgetpu,
+    torch2imx,
+    torch2onnx,
+)
+from ultralytics.utils.files import file_size
 from ultralytics.utils.metrics import batch_probiou
 from ultralytics.utils.nms import TorchNMS
 from ultralytics.utils.ops import Profile
@@ -206,15 +214,6 @@ def validate_args(format, passed_args, valid_args):
             assert arg in valid_args, f"ERROR ❌️ argument '{arg}' is not supported for format='{format}'"
 
 
-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"))
-
-
 def try_export(inner_func):
     """YOLO export decorator, i.e. @try_export."""
     inner_args = get_default_args(inner_func)
@@ -371,7 +370,7 @@ class Exporter:
                 LOGGER.warning("IMX export requires nms=True, setting nms=True.")
                 self.args.nms = True
             if model.task not in {"detect", "pose", "classify"}:
-                raise ValueError("IMX export only supported for detection and pose estimation models.")
+                raise ValueError("IMX export only supported for detection, pose estimation, and classification models.")
         if not hasattr(model, "names"):
             model.names = default_class_names()
         model.names = check_class_names(model.names)
@@ -461,6 +460,10 @@ class Exporter:
             from ultralytics.utils.export.imx import FXModel
 
             model = FXModel(model, self.imgsz)
+        if tflite or edgetpu:
+            from ultralytics.utils.export.tensorflow import tf_wrapper
+
+            model = tf_wrapper(model)
         for m in model.modules():
             if isinstance(m, Classify):
                 m.export = True
@@ -642,7 +645,7 @@ class Exporter:
             assert TORCH_1_13, f"'nms=True' ONNX export requires torch>=1.13 (found torch=={TORCH_VERSION})"
 
         f = str(self.file.with_suffix(".onnx"))
-        output_names = ["output0", "output1"] if isinstance(self.model, SegmentationModel) else ["output0"]
+        output_names = ["output0", "output1"] if self.model.task == "segment" else ["output0"]
         dynamic = self.args.dynamic
         if dynamic:
             dynamic = {"images": {0: "batch", 2: "height", 3: "width"}}  # shape(1,3,640,640)
@@ -1053,75 +1056,43 @@ class Exporter:
         if f.is_dir():
             shutil.rmtree(f)  # delete output folder
 
-        # 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)
+        # Export to TF
+        images = None
+        if self.args.int8 and self.args.data:
+            images = [batch["img"] for batch in self.get_int8_calibration_dataloader(prefix)]
+            images = (
+                torch.nn.functional.interpolate(torch.cat(images, 0).float(), size=self.imgsz)
+                .permute(0, 2, 3, 1)
+                .numpy()
+                .astype(np.float32)
+            )
 
         # Export to ONNX
         if isinstance(self.model.model[-1], RTDETRDecoder):
             self.args.opset = self.args.opset or 19
             assert 16 <= self.args.opset <= 19, "RTDETR export requires opset>=16;<=19"
         self.args.simplify = True
-        f_onnx = self.export_onnx()
-
-        # Export to TF
-        np_data = None
-        if self.args.int8:
-            tmp_file = f / "tmp_tflite_int8_calibration_images.npy"  # int8 calibration images file
-            if self.args.data:
-                f.mkdir()
-                images = [batch["img"] for batch in self.get_int8_calibration_dataloader(prefix)]
-                images = torch.nn.functional.interpolate(torch.cat(images, 0).float(), size=self.imgsz).permute(
-                    0, 2, 3, 1
-                )
-                np.save(str(tmp_file), images.numpy().astype(np.float32))  # 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=f_onnx,
-            output_folder_path=str(f),
-            not_use_onnxsim=True,
-            verbosity="error",  # note INT8-FP16 activation bug https://github.com/ultralytics/ultralytics/issues/15873
-            output_integer_quantized_tflite=self.args.int8,
-            custom_input_op_name_np_data_path=np_data,
-            enable_batchmatmul_unfold=True and not self.args.int8,  # fix lower no. of detected objects on GPU delegate
-            output_signaturedefs=True,  # fix error with Attention block group convolution
-            disable_group_convolution=self.args.format in {"tfjs", "edgetpu"},  # fix error with group convolution
+        f_onnx = self.export_onnx()  # ensure ONNX is available
+        keras_model = onnx2saved_model(
+            f_onnx,
+            f,
+            int8=self.args.int8,
+            images=images,
+            disable_group_convolution=self.args.format in {"tfjs", "edgetpu"},
+            prefix=prefix,
         )
         YAML.save(f / "metadata.yaml", self.metadata)  # add metadata.yaml
-
-        # Remove/rename TFLite models
-        if self.args.int8:
-            tmp_file.unlink(missing_ok=True)
-            for file in f.rglob("*_dynamic_range_quant.tflite"):
-                file.rename(file.with_name(file.stem.replace("_dynamic_range_quant", "_int8") + file.suffix))
-            for file in f.rglob("*_integer_quant_with_int16_act.tflite"):
-                file.unlink()  # delete extra fp16 activation TFLite files
-
         # Add TFLite metadata
         for file in f.rglob("*.tflite"):
-            f.unlink() if "quant_with_int16_act.tflite" in str(f) else self._add_tflite_metadata(file)
+            file.unlink() if "quant_with_int16_act.tflite" in str(file) else self._add_tflite_metadata(file)
 
         return str(f), keras_model  # or keras_model = tf.saved_model.load(f, tags=None, options=None)
 
     @try_export
     def export_pb(self, keras_model, prefix=colorstr("TensorFlow GraphDef:")):
         """Export YOLO model to TensorFlow GraphDef *.pb format https://github.com/leimao/Frozen-Graph-TensorFlow."""
-        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__}...")
         f = self.file.with_suffix(".pb")
-
-        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(f.parent), name=f.name, as_text=False)
+        keras2pb(keras_model, f, prefix)
         return f
 
     @try_export
@@ -1189,22 +1160,11 @@ class Exporter:
                 "sudo apt-get install edgetpu-compiler",
             ):
                 subprocess.run(c if is_sudo_available() else c.replace("sudo ", ""), shell=True, check=True)
-        ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().rsplit(maxsplit=1)[-1]
 
+        ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().rsplit(maxsplit=1)[-1]
         LOGGER.info(f"\n{prefix} starting export with Edge TPU compiler {ver}...")
+        tflite2edgetpu(tflite_file=tflite_model, output_dir=tflite_model.parent, prefix=prefix)
         f = str(tflite_model).replace(".tflite", "_edgetpu.tflite")  # Edge TPU model
-
-        cmd = (
-            "edgetpu_compiler "
-            f'--out_dir "{Path(f).parent}" '
-            "--show_operations "
-            "--search_delegate "
-            "--delegate_search_step 30 "
-            "--timeout_sec 180 "
-            f'"{tflite_model}"'
-        )
-        LOGGER.info(f"{prefix} running '{cmd}'")
-        subprocess.run(cmd, shell=True)
         self._add_tflite_metadata(f)
         return f
 
@@ -1212,31 +1172,10 @@ class Exporter:
     def export_tfjs(self, prefix=colorstr("TensorFlow.js:")):
         """Export YOLO model to TensorFlow.js format."""
         check_requirements("tensorflowjs")
-        import tensorflow as tf
-        import tensorflowjs as tfjs
 
-        LOGGER.info(f"\n{prefix} starting export with tensorflowjs {tfjs.__version__}...")
         f = str(self.file).replace(self.file.suffix, "_web_model")  # js dir
         f_pb = str(self.file.with_suffix(".pb"))  # *.pb path
-
-        gd = tf.Graph().as_graph_def()  # TF GraphDef
-        with open(f_pb, "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 self.args.half else "--quantize_uint8" if self.args.int8 else ""
-        with spaces_in_path(f_pb) as fpb_, spaces_in_path(f) 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 f:
-            LOGGER.warning(f"{prefix} your model may not work correctly with spaces in path '{f}'.")
-
+        pb2tfjs(pb_file=f_pb, output_dir=f, half=self.args.half, int8=self.args.int8, prefix=prefix)
         # Add metadata
         YAML.save(Path(f) / "metadata.yaml", self.metadata)  # add metadata.yaml
         return f

ultralytics/models/rtdetr/val.py

@@ -89,7 +89,7 @@ class RTDETRDataset(YOLODataset):
             transforms = v8_transforms(self, self.imgsz, hyp, stretch=True)
         else:
             # transforms = Compose([LetterBox(new_shape=(self.imgsz, self.imgsz), auto=False, scale_fill=True)])
-            transforms = Compose([])
+            transforms = Compose([lambda x: {**x, **{"ratio_pad": [x["ratio_pad"], [0, 0]]}}])
         transforms.append(
             Format(
                 bbox_format="xywh",

ultralytics/nn/autobackend.py

@@ -428,7 +428,7 @@ class AutoBackend(nn.Module):
             LOGGER.info(f"Loading {w} for TensorFlow GraphDef inference...")
             import tensorflow as tf
 
-            from ultralytics.engine.exporter import gd_outputs
+            from ultralytics.utils.export.tensorflow import gd_outputs
 
             def wrap_frozen_graph(gd, inputs, outputs):
                 """Wrap frozen graphs for deployment."""

ultralytics/nn/modules/head.py

@@ -166,22 +166,8 @@ class Detect(nn.Module):
             self.anchors, self.strides = (x.transpose(0, 1) for x in make_anchors(x, self.stride, 0.5))
             self.shape = shape
 
-        if self.export and self.format in {"saved_model", "pb", "tflite", "edgetpu", "tfjs"}:  # avoid TF FlexSplitV ops
-            box = x_cat[:, : self.reg_max * 4]
-            cls = x_cat[:, self.reg_max * 4 :]
-        else:
-            box, cls = x_cat.split((self.reg_max * 4, self.nc), 1)
-
-        if self.export and self.format in {"tflite", "edgetpu"}:
-            # Precompute normalization factor to increase numerical stability
-            # See https://github.com/ultralytics/ultralytics/issues/7371
-            grid_h = shape[2]
-            grid_w = 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])
-        else:
-            dbox = self.decode_bboxes(self.dfl(box), self.anchors.unsqueeze(0)) * self.strides
+        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 torch.cat((dbox, cls.sigmoid()), 1)
 
     def bias_init(self):
@@ -391,20 +377,9 @@ class Pose(Detect):
         """Decode keypoints from predictions."""
         ndim = self.kpt_shape[1]
         if self.export:
-            if self.format in {
-                "tflite",
-                "edgetpu",
-            }:  # 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
-            else:
-                # NCNN fix
-                y = kpts.view(bs, *self.kpt_shape, -1)
-                a = (y[:, :, :2] * 2.0 + (self.anchors - 0.5)) * self.strides
+            # NCNN fix
+            y = kpts.view(bs, *self.kpt_shape, -1)
+            a = (y[:, :, :2] * 2.0 + (self.anchors - 0.5)) * self.strides
             if ndim == 3:
                 a = torch.cat((a, y[:, :, 2:3].sigmoid()), 2)
             return a.view(bs, self.nk, -1)

ultralytics/utils/export/__init__.py

@@ -1,242 +1,7 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
-from __future__ import annotations
+from .engine import onnx2engine, torch2onnx
+from .imx import torch2imx
+from .tensorflow import keras2pb, onnx2saved_model, pb2tfjs, tflite2edgetpu
 
-import json
-from pathlib import Path
-
-import torch
-
-from ultralytics.utils import IS_JETSON, LOGGER
-from ultralytics.utils.torch_utils import TORCH_2_4
-
-from .imx import torch2imx  # noqa
-
-
-def torch2onnx(
-    torch_model: torch.nn.Module,
-    im: torch.Tensor,
-    onnx_file: str,
-    opset: int = 14,
-    input_names: list[str] = ["images"],
-    output_names: list[str] = ["output0"],
-    dynamic: bool | dict = False,
-) -> None:
-    """
-    Export a PyTorch model to ONNX format.
-
-    Args:
-        torch_model (torch.nn.Module): The PyTorch model to export.
-        im (torch.Tensor): Example input tensor for the model.
-        onnx_file (str): Path to save the exported ONNX file.
-        opset (int): ONNX opset version to use for export.
-        input_names (list[str]): List of input tensor names.
-        output_names (list[str]): List of output tensor names.
-        dynamic (bool | dict, optional): Whether to enable dynamic axes.
-
-    Notes:
-        Setting `do_constant_folding=True` may cause issues with DNN inference for torch>=1.12.
-    """
-    kwargs = {"dynamo": False} if TORCH_2_4 else {}
-    torch.onnx.export(
-        torch_model,
-        im,
-        onnx_file,
-        verbose=False,
-        opset_version=opset,
-        do_constant_folding=True,  # WARNING: DNN inference with torch>=1.12 may require do_constant_folding=False
-        input_names=input_names,
-        output_names=output_names,
-        dynamic_axes=dynamic or None,
-        **kwargs,
-    )
-
-
-def onnx2engine(
-    onnx_file: str,
-    engine_file: str | None = None,
-    workspace: int | None = None,
-    half: bool = False,
-    int8: bool = False,
-    dynamic: bool = False,
-    shape: tuple[int, int, int, int] = (1, 3, 640, 640),
-    dla: int | None = None,
-    dataset=None,
-    metadata: dict | None = None,
-    verbose: bool = False,
-    prefix: str = "",
-) -> None:
-    """
-    Export a YOLO model to TensorRT engine format.
-
-    Args:
-        onnx_file (str): Path to the ONNX file to be converted.
-        engine_file (str, optional): Path to save the generated TensorRT engine file.
-        workspace (int, optional): Workspace size in GB for TensorRT.
-        half (bool, optional): Enable FP16 precision.
-        int8 (bool, optional): Enable INT8 precision.
-        dynamic (bool, optional): Enable dynamic input shapes.
-        shape (tuple[int, int, int, int], optional): Input shape (batch, channels, height, width).
-        dla (int, optional): DLA core to use (Jetson devices only).
-        dataset (ultralytics.data.build.InfiniteDataLoader, optional): Dataset for INT8 calibration.
-        metadata (dict, optional): Metadata to include in the engine file.
-        verbose (bool, optional): Enable verbose logging.
-        prefix (str, optional): Prefix for log messages.
-
-    Raises:
-        ValueError: If DLA is enabled on non-Jetson devices or required precision is not set.
-        RuntimeError: If the ONNX file cannot be parsed.
-
-    Notes:
-        TensorRT version compatibility is handled for workspace size and engine building.
-        INT8 calibration requires a dataset and generates a calibration cache.
-        Metadata is serialized and written to the engine file if provided.
-    """
-    import tensorrt as trt
-
-    engine_file = engine_file or Path(onnx_file).with_suffix(".engine")
-
-    logger = trt.Logger(trt.Logger.INFO)
-    if verbose:
-        logger.min_severity = trt.Logger.Severity.VERBOSE
-
-    # Engine builder
-    builder = trt.Builder(logger)
-    config = builder.create_builder_config()
-    workspace_bytes = int((workspace or 0) * (1 << 30))
-    is_trt10 = int(trt.__version__.split(".", 1)[0]) >= 10  # is TensorRT >= 10
-    if is_trt10 and workspace_bytes > 0:
-        config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace_bytes)
-    elif workspace_bytes > 0:  # TensorRT versions 7, 8
-        config.max_workspace_size = workspace_bytes
-    flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
-    network = builder.create_network(flag)
-    half = builder.platform_has_fast_fp16 and half
-    int8 = builder.platform_has_fast_int8 and int8
-
-    # Optionally switch to DLA if enabled
-    if dla is not None:
-        if not IS_JETSON:
-            raise ValueError("DLA is only available on NVIDIA Jetson devices")
-        LOGGER.info(f"{prefix} enabling DLA on core {dla}...")
-        if not half and not int8:
-            raise ValueError(
-                "DLA requires either 'half=True' (FP16) or 'int8=True' (INT8) to be enabled. Please enable one of them and try again."
-            )
-        config.default_device_type = trt.DeviceType.DLA
-        config.DLA_core = int(dla)
-        config.set_flag(trt.BuilderFlag.GPU_FALLBACK)
-
-    # Read ONNX file
-    parser = trt.OnnxParser(network, logger)
-    if not parser.parse_from_file(onnx_file):
-        raise RuntimeError(f"failed to load ONNX file: {onnx_file}")
-
-    # Network inputs
-    inputs = [network.get_input(i) for i in range(network.num_inputs)]
-    outputs = [network.get_output(i) for i in range(network.num_outputs)]
-    for inp in inputs:
-        LOGGER.info(f'{prefix} input "{inp.name}" with shape{inp.shape} {inp.dtype}')
-    for out in outputs:
-        LOGGER.info(f'{prefix} output "{out.name}" with shape{out.shape} {out.dtype}')
-
-    if dynamic:
-        profile = builder.create_optimization_profile()
-        min_shape = (1, shape[1], 32, 32)  # minimum input shape
-        max_shape = (*shape[:2], *(int(max(2, workspace or 2) * d) for d in shape[2:]))  # max input shape
-        for inp in inputs:
-            profile.set_shape(inp.name, min=min_shape, opt=shape, max=max_shape)
-        config.add_optimization_profile(profile)
-        if int8:
-            config.set_calibration_profile(profile)
-
-    LOGGER.info(f"{prefix} building {'INT8' if int8 else 'FP' + ('16' if half else '32')} engine as {engine_file}")
-    if int8:
-        config.set_flag(trt.BuilderFlag.INT8)
-        config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED
-
-        class EngineCalibrator(trt.IInt8Calibrator):
-            """
-            Custom INT8 calibrator for TensorRT engine optimization.
-
-            This calibrator provides the necessary interface for TensorRT to perform INT8 quantization calibration
-            using a dataset. It handles batch generation, caching, and calibration algorithm selection.
-
-            Attributes:
-                dataset: Dataset for calibration.
-                data_iter: Iterator over the calibration dataset.
-                algo (trt.CalibrationAlgoType): Calibration algorithm type.
-                batch (int): Batch size for calibration.
-                cache (Path): Path to save the calibration cache.
-
-            Methods:
-                get_algorithm: Get the calibration algorithm to use.
-                get_batch_size: Get the batch size to use for calibration.
-                get_batch: Get the next batch to use for calibration.
-                read_calibration_cache: Use existing cache instead of calibrating again.
-                write_calibration_cache: Write calibration cache to disk.
-            """
-
-            def __init__(
-                self,
-                dataset,  # ultralytics.data.build.InfiniteDataLoader
-                cache: str = "",
-            ) -> None:
-                """Initialize the INT8 calibrator with dataset and cache path."""
-                trt.IInt8Calibrator.__init__(self)
-                self.dataset = dataset
-                self.data_iter = iter(dataset)
-                self.algo = (
-                    trt.CalibrationAlgoType.ENTROPY_CALIBRATION_2  # DLA quantization needs ENTROPY_CALIBRATION_2
-                    if dla is not None
-                    else trt.CalibrationAlgoType.MINMAX_CALIBRATION
-                )
-                self.batch = dataset.batch_size
-                self.cache = Path(cache)
-
-            def get_algorithm(self) -> trt.CalibrationAlgoType:
-                """Get the calibration algorithm to use."""
-                return self.algo
-
-            def get_batch_size(self) -> int:
-                """Get the batch size to use for calibration."""
-                return self.batch or 1
-
-            def get_batch(self, names) -> list[int] | None:
-                """Get the next batch to use for calibration, as a list of device memory pointers."""
-                try:
-                    im0s = next(self.data_iter)["img"] / 255.0
-                    im0s = im0s.to("cuda") if im0s.device.type == "cpu" else im0s
-                    return [int(im0s.data_ptr())]
-                except StopIteration:
-                    # Return None to signal to TensorRT there is no calibration data remaining
-                    return None
-
-            def read_calibration_cache(self) -> bytes | None:
-                """Use existing cache instead of calibrating again, otherwise, implicitly return None."""
-                if self.cache.exists() and self.cache.suffix == ".cache":
-                    return self.cache.read_bytes()
-
-            def write_calibration_cache(self, cache: bytes) -> None:
-                """Write calibration cache to disk."""
-                _ = self.cache.write_bytes(cache)
-
-        # Load dataset w/ builder (for batching) and calibrate
-        config.int8_calibrator = EngineCalibrator(
-            dataset=dataset,
-            cache=str(Path(onnx_file).with_suffix(".cache")),
-        )
-
-    elif half:
-        config.set_flag(trt.BuilderFlag.FP16)
-
-    # Write file
-    build = builder.build_serialized_network if is_trt10 else builder.build_engine
-    with build(network, config) as engine, open(engine_file, "wb") as t:
-        # Metadata
-        if metadata is not None:
-            meta = json.dumps(metadata)
-            t.write(len(meta).to_bytes(4, byteorder="little", signed=True))
-            t.write(meta.encode())
-        # Model
-        t.write(engine if is_trt10 else engine.serialize())
+__all__ = ["keras2pb", "onnx2engine", "onnx2saved_model", "pb2tfjs", "tflite2edgetpu", "torch2imx", "torch2onnx"]

ultralytics/utils/export/engine.py

@@ -0,0 +1,240 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+
+from ultralytics.utils import IS_JETSON, LOGGER
+from ultralytics.utils.torch_utils import TORCH_2_4
+
+
+def torch2onnx(
+    torch_model: torch.nn.Module,
+    im: torch.Tensor,
+    onnx_file: str,
+    opset: int = 14,
+    input_names: list[str] = ["images"],
+    output_names: list[str] = ["output0"],
+    dynamic: bool | dict = False,
+) -> None:
+    """
+    Export a PyTorch model to ONNX format.
+
+    Args:
+        torch_model (torch.nn.Module): The PyTorch model to export.
+        im (torch.Tensor): Example input tensor for the model.
+        onnx_file (str): Path to save the exported ONNX file.
+        opset (int): ONNX opset version to use for export.
+        input_names (list[str]): List of input tensor names.
+        output_names (list[str]): List of output tensor names.
+        dynamic (bool | dict, optional): Whether to enable dynamic axes.
+
+    Notes:
+        Setting `do_constant_folding=True` may cause issues with DNN inference for torch>=1.12.
+    """
+    kwargs = {"dynamo": False} if TORCH_2_4 else {}
+    torch.onnx.export(
+        torch_model,
+        im,
+        onnx_file,
+        verbose=False,
+        opset_version=opset,
+        do_constant_folding=True,  # WARNING: DNN inference with torch>=1.12 may require do_constant_folding=False
+        input_names=input_names,
+        output_names=output_names,
+        dynamic_axes=dynamic or None,
+        **kwargs,
+    )
+
+
+def onnx2engine(
+    onnx_file: str,
+    engine_file: str | None = None,
+    workspace: int | None = None,
+    half: bool = False,
+    int8: bool = False,
+    dynamic: bool = False,
+    shape: tuple[int, int, int, int] = (1, 3, 640, 640),
+    dla: int | None = None,
+    dataset=None,
+    metadata: dict | None = None,
+    verbose: bool = False,
+    prefix: str = "",
+) -> None:
+    """
+    Export a YOLO model to TensorRT engine format.
+
+    Args:
+        onnx_file (str): Path to the ONNX file to be converted.
+        engine_file (str, optional): Path to save the generated TensorRT engine file.
+        workspace (int, optional): Workspace size in GB for TensorRT.
+        half (bool, optional): Enable FP16 precision.
+        int8 (bool, optional): Enable INT8 precision.
+        dynamic (bool, optional): Enable dynamic input shapes.
+        shape (tuple[int, int, int, int], optional): Input shape (batch, channels, height, width).
+        dla (int, optional): DLA core to use (Jetson devices only).
+        dataset (ultralytics.data.build.InfiniteDataLoader, optional): Dataset for INT8 calibration.
+        metadata (dict, optional): Metadata to include in the engine file.
+        verbose (bool, optional): Enable verbose logging.
+        prefix (str, optional): Prefix for log messages.
+
+    Raises:
+        ValueError: If DLA is enabled on non-Jetson devices or required precision is not set.
+        RuntimeError: If the ONNX file cannot be parsed.
+
+    Notes:
+        TensorRT version compatibility is handled for workspace size and engine building.
+        INT8 calibration requires a dataset and generates a calibration cache.
+        Metadata is serialized and written to the engine file if provided.
+    """
+    import tensorrt as trt
+
+    engine_file = engine_file or Path(onnx_file).with_suffix(".engine")
+
+    logger = trt.Logger(trt.Logger.INFO)
+    if verbose:
+        logger.min_severity = trt.Logger.Severity.VERBOSE
+
+    # Engine builder
+    builder = trt.Builder(logger)
+    config = builder.create_builder_config()
+    workspace_bytes = int((workspace or 0) * (1 << 30))
+    is_trt10 = int(trt.__version__.split(".", 1)[0]) >= 10  # is TensorRT >= 10
+    if is_trt10 and workspace_bytes > 0:
+        config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace_bytes)
+    elif workspace_bytes > 0:  # TensorRT versions 7, 8
+        config.max_workspace_size = workspace_bytes
+    flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
+    network = builder.create_network(flag)
+    half = builder.platform_has_fast_fp16 and half
+    int8 = builder.platform_has_fast_int8 and int8
+
+    # Optionally switch to DLA if enabled
+    if dla is not None:
+        if not IS_JETSON:
+            raise ValueError("DLA is only available on NVIDIA Jetson devices")
+        LOGGER.info(f"{prefix} enabling DLA on core {dla}...")
+        if not half and not int8:
+            raise ValueError(
+                "DLA requires either 'half=True' (FP16) or 'int8=True' (INT8) to be enabled. Please enable one of them and try again."
+            )
+        config.default_device_type = trt.DeviceType.DLA
+        config.DLA_core = int(dla)
+        config.set_flag(trt.BuilderFlag.GPU_FALLBACK)
+
+    # Read ONNX file
+    parser = trt.OnnxParser(network, logger)
+    if not parser.parse_from_file(onnx_file):
+        raise RuntimeError(f"failed to load ONNX file: {onnx_file}")
+
+    # Network inputs
+    inputs = [network.get_input(i) for i in range(network.num_inputs)]
+    outputs = [network.get_output(i) for i in range(network.num_outputs)]
+    for inp in inputs:
+        LOGGER.info(f'{prefix} input "{inp.name}" with shape{inp.shape} {inp.dtype}')
+    for out in outputs:
+        LOGGER.info(f'{prefix} output "{out.name}" with shape{out.shape} {out.dtype}')
+
+    if dynamic:
+        profile = builder.create_optimization_profile()
+        min_shape = (1, shape[1], 32, 32)  # minimum input shape
+        max_shape = (*shape[:2], *(int(max(2, workspace or 2) * d) for d in shape[2:]))  # max input shape
+        for inp in inputs:
+            profile.set_shape(inp.name, min=min_shape, opt=shape, max=max_shape)
+        config.add_optimization_profile(profile)
+        if int8:
+            config.set_calibration_profile(profile)
+
+    LOGGER.info(f"{prefix} building {'INT8' if int8 else 'FP' + ('16' if half else '32')} engine as {engine_file}")
+    if int8:
+        config.set_flag(trt.BuilderFlag.INT8)
+        config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED
+
+        class EngineCalibrator(trt.IInt8Calibrator):
+            """
+            Custom INT8 calibrator for TensorRT engine optimization.
+
+            This calibrator provides the necessary interface for TensorRT to perform INT8 quantization calibration
+            using a dataset. It handles batch generation, caching, and calibration algorithm selection.
+
+            Attributes:
+                dataset: Dataset for calibration.
+                data_iter: Iterator over the calibration dataset.
+                algo (trt.CalibrationAlgoType): Calibration algorithm type.
+                batch (int): Batch size for calibration.
+                cache (Path): Path to save the calibration cache.
+
+            Methods:
+                get_algorithm: Get the calibration algorithm to use.
+                get_batch_size: Get the batch size to use for calibration.
+                get_batch: Get the next batch to use for calibration.
+                read_calibration_cache: Use existing cache instead of calibrating again.
+                write_calibration_cache: Write calibration cache to disk.
+            """
+
+            def __init__(
+                self,
+                dataset,  # ultralytics.data.build.InfiniteDataLoader
+                cache: str = "",
+            ) -> None:
+                """Initialize the INT8 calibrator with dataset and cache path."""
+                trt.IInt8Calibrator.__init__(self)
+                self.dataset = dataset
+                self.data_iter = iter(dataset)
+                self.algo = (
+                    trt.CalibrationAlgoType.ENTROPY_CALIBRATION_2  # DLA quantization needs ENTROPY_CALIBRATION_2
+                    if dla is not None
+                    else trt.CalibrationAlgoType.MINMAX_CALIBRATION
+                )
+                self.batch = dataset.batch_size
+                self.cache = Path(cache)
+
+            def get_algorithm(self) -> trt.CalibrationAlgoType:
+                """Get the calibration algorithm to use."""
+                return self.algo
+
+            def get_batch_size(self) -> int:
+                """Get the batch size to use for calibration."""
+                return self.batch or 1
+
+            def get_batch(self, names) -> list[int] | None:
+                """Get the next batch to use for calibration, as a list of device memory pointers."""
+                try:
+                    im0s = next(self.data_iter)["img"] / 255.0
+                    im0s = im0s.to("cuda") if im0s.device.type == "cpu" else im0s
+                    return [int(im0s.data_ptr())]
+                except StopIteration:
+                    # Return None to signal to TensorRT there is no calibration data remaining
+                    return None
+
+            def read_calibration_cache(self) -> bytes | None:
+                """Use existing cache instead of calibrating again, otherwise, implicitly return None."""
+                if self.cache.exists() and self.cache.suffix == ".cache":
+                    return self.cache.read_bytes()
+
+            def write_calibration_cache(self, cache: bytes) -> None:
+                """Write calibration cache to disk."""
+                _ = self.cache.write_bytes(cache)
+
+        # Load dataset w/ builder (for batching) and calibrate
+        config.int8_calibrator = EngineCalibrator(
+            dataset=dataset,
+            cache=str(Path(onnx_file).with_suffix(".cache")),
+        )
+
+    elif half:
+        config.set_flag(trt.BuilderFlag.FP16)
+
+    # Write file
+    build = builder.build_serialized_network if is_trt10 else builder.build_engine
+    with build(network, config) as engine, open(engine_file, "wb") as t:
+        # Metadata
+        if metadata is not None:
+            meta = json.dumps(metadata)
+            t.write(len(meta).to_bytes(4, byteorder="little", signed=True))
+            t.write(meta.encode())
+        # Model
+        t.write(engine if is_trt10 else engine.serialize())

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/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