ADAS/HSAP
2
2
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: initial HSAP platform

Browse Source 
Huaxu Sentinel Active Safety Platform with embedded algorithm code,
Docker Compose setup, and vendored dataset scaffolds for clone-and-run.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Chengfang Lu
2026-05-25 16:59:59 +08:00
commit 7c43b44c57
1619 changed files with 373355 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
algorithms/dms_yolo/code/ultralytics/utils/export/__init__.py Normal file
View File

@@ -0,0 +1,7 @@
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from .engine import onnx2engine, torch2onnx
from .imx import torch2imx
from .tensorflow import keras2pb, onnx2saved_model, pb2tfjs, tflite2edgetpu
__all__ = ["keras2pb", "onnx2engine", "onnx2saved_model", "pb2tfjs", "tflite2edgetpu", "torch2imx", "torch2onnx"]

246
algorithms/dms_yolo/code/ultralytics/utils/export/engine.py Normal file
View File

@@ -0,0 +1,246 @@
# 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 and not is_trt10: # deprecated in TensorRT 10, causes internal errors
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
if is_trt10:
# TensorRT 10+ returns bytes directly, not a context manager
engine = builder.build_serialized_network(network, config)
if engine is None:
raise RuntimeError("TensorRT engine build failed, check logs for errors")
with open(engine_file, "wb") as t:
if metadata is not None:
meta = json.dumps(metadata)
t.write(len(meta).to_bytes(4, byteorder="little", signed=True))
t.write(meta.encode())
t.write(engine)
else:
with builder.build_engine(network, config) as engine, open(engine_file, "wb") as t:
if metadata is not None:
meta = json.dumps(metadata)
t.write(len(meta).to_bytes(4, byteorder="little", signed=True))
t.write(meta.encode())
t.write(engine.serialize())

343
algorithms/dms_yolo/code/ultralytics/utils/export/imx.py Normal file
View File

@@ -0,0 +1,343 @@
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
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, 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_19"],
"weights_memory": 2585350.2439,
"n_layers": {238, 239},
},
"pose": {
"layer_names": ["sub", "mul_2", "add_14", "cat_21", "cat_22", "mul_4", "add_15"],
"weights_memory": 2437771.67,
"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},
},
},
"YOLOv8": {
"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_17", "mul", "sub", "add_6", "cat_18"],
"weights_memory": 2482451.85,
"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},
},
},
}
class FXModel(torch.nn.Module):
"""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
copying.
Attributes:
model (nn.Module): The original model's layers.
"""
def __init__(self, model, imgsz=(640, 640)):
"""Initialize the FXModel.
Args:
model (nn.Module): The original model to wrap for torch.fx compatibility.
imgsz (tuple[int, int]): The input image size (height, width). Default is (640, 640).
"""
super().__init__()
copy_attr(self, model)
# Explicitly set `model` since `copy_attr` somehow does not copy it.
self.model = model.model
self.imgsz = imgsz
def forward(self, x):
"""Forward pass through the model.
This method performs the forward pass through the model, handling the dependencies between layers and saving
intermediate outputs.
Args:
x (torch.Tensor): The input tensor to the model.
Returns:
(torch.Tensor): The output tensor from the model.
"""
y = [] # outputs
for m in self.model:
if m.f != -1: # if not from previous layer
# from earlier layers
x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f]
if isinstance(m, Detect):
m._inference = types.MethodType(_inference, m) # bind method to Detect
m.anchors, m.strides = (
x.transpose(0, 1)
for x in make_anchors(
torch.cat([s / m.stride.unsqueeze(-1) for s in self.imgsz], dim=1), m.stride, 0.5
)
)
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
def _inference(self, x: dict[str, torch.Tensor]) -> tuple[torch.Tensor]:
"""Decode boxes and cls scores for imx object detection."""
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
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(kpt)
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 edge-mdt-cl."""
def __init__(
self,
model: torch.nn.Module,
score_threshold: float = 0.001,
iou_threshold: float = 0.7,
max_detections: int = 300,
task: str = "detect",
):
"""Initialize NMSWrapper with PyTorch Module and NMS parameters.
Args:
model (torch.nn.Module): Model instance.
score_threshold (float): Score threshold for non-maximum suppression.
iou_threshold (float): Intersection over union threshold for non-maximum suppression.
max_detections (int): The number of detections to return.
task (str): Task type, either 'detect' or 'pose'.
"""
super().__init__()
self.model = model
self.score_threshold = score_threshold
self.iou_threshold = iou_threshold
self.max_detections = max_detections
self.task = task
def forward(self, images):
"""Forward pass with model inference and NMS post-processing."""
from edgemdt_cl.pytorch.nms.nms_with_indices import multiclass_nms_with_indices
# model inference
outputs = self.model(images)
boxes, scores = outputs[0], outputs[1]
nms_outputs = multiclass_nms_with_indices(
boxes=boxes,
scores=scores,
score_threshold=self.score_threshold,
iou_threshold=self.iou_threshold,
max_detections=self.max_detections,
)
if self.task == "pose":
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
def torch2imx(
model: torch.nn.Module,
file: Path | str,
conf: float,
iou: float,
max_det: int,
metadata: dict | None = None,
gptq: bool = False,
dataset=None,
prefix: str = "",
):
"""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.
Args:
model (torch.nn.Module): The YOLO model to export. Must be YOLOv8n or YOLO11n.
file (Path | str): Output file path for the exported model.
conf (float): Confidence threshold for NMS post-processing.
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.
dataset (optional): Representative dataset for quantization calibration. Defaults to None.
prefix (str, optional): Logging prefix string. Defaults to "".
Returns:
f (Path): Path to the exported IMX model directory
Raises:
ValueError: If the model is not a supported YOLOv8n or YOLO11n variant.
Examples:
>>> from ultralytics import YOLO
>>> model = YOLO("yolo11n.pt")
>>> 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
- Only supports YOLOv8n and YOLO11n models (detection and pose tasks)
- Output includes quantized ONNX model, IMX binary, and labels.txt file
"""
import model_compression_toolkit as mct
import onnx
from edgemdt_tpc import get_target_platform_capabilities
LOGGER.info(f"\n{prefix} starting export with model_compression_toolkit {mct.__version__}...")
def representative_dataset_gen(dataloader=dataset):
for batch in dataloader:
img = batch["img"]
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()
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())) 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"]:
bit_cfg.set_manual_activation_bit_width([mct.core.common.network_editors.NodeNameFilter(layer_name)], 16)
config = mct.core.CoreConfig(
mixed_precision_config=mct.core.MixedPrecisionQuantizationConfig(num_of_images=10),
quantization_config=mct.core.QuantizationConfig(concat_threshold_update=True),
bit_width_config=bit_cfg,
)
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
model=model,
representative_data_gen=representative_dataset_gen,
target_resource_utilization=resource_utilization,
gptq_config=mct.gptq.get_pytorch_gptq_config(
n_epochs=1000, use_hessian_based_weights=False, use_hessian_sample_attention=False
),
core_config=config,
target_platform_capabilities=tpc,
)[0]
if gptq
else mct.ptq.pytorch_post_training_quantization( # Perform post training quantization
in_module=model,
representative_data_gen=representative_dataset_gen,
target_resource_utilization=resource_utilization,
core_config=config,
target_platform_capabilities=tpc,
)[0]
)
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
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():
meta = model_onnx.metadata_props.add()
meta.key, meta.value = k, str(v)
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(
[str(imxconv), "-i", str(onnx_model), "-o", str(f), "--no-input-persistency", "--overwrite-output"],
check=True,
)
# Needed for imx models.
with open(f / "labels.txt", "w", encoding="utf-8") as file:
file.writelines([f"{name}\n" for _, name in model.names.items()])
return f

231
algorithms/dms_yolo/code/ultralytics/utils/export/tensorflow.py Normal file
View File

@@ -0,0 +1,231 @@
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
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, Pose26
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 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._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_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: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(boxes) * norm, self.anchors.unsqueeze(0) * norm[:, :2])
return dbox
def _tf_kpts_decode(self, kpts: torch.Tensor, is_pose26: bool = False) -> torch.Tensor:
"""Decode keypoints for TensorFlow export."""
ndim = self.kpt_shape[1]
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: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] + 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)
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"))