409 lines
17 KiB
Python
409 lines
17 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import json
|
||
|
|
import sys
|
||
|
|
from copy import deepcopy
|
||
|
|
from pathlib import Path
|
||
|
|
from types import SimpleNamespace
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
import torch
|
||
|
|
import torch.nn as nn
|
||
|
|
|
||
|
|
FILE = Path(__file__).resolve()
|
||
|
|
ROOT = FILE.parents[2]
|
||
|
|
if str(ROOT) not in sys.path:
|
||
|
|
sys.path.append(str(ROOT))
|
||
|
|
|
||
|
|
from ultralytics.nn.modules import Detect, RTDETRDecoder
|
||
|
|
from ultralytics.nn.tasks import load_checkpoint
|
||
|
|
from ultralytics.utils import LOGGER
|
||
|
|
from ultralytics.utils.patches import arange_patch, onnx_export_patch
|
||
|
|
|
||
|
|
|
||
|
|
DEFAULT_MODEL = ROOT / "runs" / "detect" / "train_mono3d_roi0_202603291330_epoch61.pt"
|
||
|
|
DEFAULT_SAVE_DIR = ROOT / "runs" / "export" / "train_mono3d_single_roi0_202603291330"
|
||
|
|
DEFAULT_IMGSZ_WH = (768, 352)
|
||
|
|
|
||
|
|
|
||
|
|
def parse_args() -> argparse.Namespace:
|
||
|
|
parser = argparse.ArgumentParser(description="Export one yolo26 single-ROI Detect3D checkpoint.")
|
||
|
|
parser.add_argument("--model-path", type=str, default=str(DEFAULT_MODEL), help="Path to ROI checkpoint")
|
||
|
|
parser.add_argument("--save-dir", type=str, default=str(DEFAULT_SAVE_DIR), help="Directory used to store exported files")
|
||
|
|
parser.add_argument(
|
||
|
|
"--imgsz",
|
||
|
|
nargs=2,
|
||
|
|
type=int,
|
||
|
|
default=list(DEFAULT_IMGSZ_WH),
|
||
|
|
metavar=("W", "H"),
|
||
|
|
help="Example input size as width height, e.g. --imgsz 768 352",
|
||
|
|
)
|
||
|
|
parser.add_argument("--opset", type=int, default=15, help="ONNX opset version")
|
||
|
|
parser.add_argument("--max-det", type=int, default=300, help="Top-k detections kept by the Detect3D head")
|
||
|
|
parser.add_argument("--dynamic", action="store_true", help="Export with dynamic batch dimension")
|
||
|
|
parser.add_argument("--simplify", action="store_true", help="Try to simplify the exported ONNX graph with onnxslim")
|
||
|
|
export_group = parser.add_mutually_exclusive_group()
|
||
|
|
export_group.add_argument(
|
||
|
|
"--detections-only",
|
||
|
|
action="store_true",
|
||
|
|
help="Export only postprocessed 2D detections.",
|
||
|
|
)
|
||
|
|
export_group.add_argument(
|
||
|
|
"--hybrid-outputs",
|
||
|
|
action="store_true",
|
||
|
|
help="Export postprocessed 2D detections together with raw 2D/3D/edge head outputs.",
|
||
|
|
)
|
||
|
|
export_group.add_argument(
|
||
|
|
"--denorm-branch-outputs",
|
||
|
|
action="store_true",
|
||
|
|
help="Export one2one branch outputs after Detect3D branch denormalization, but before top-k postprocessing.",
|
||
|
|
)
|
||
|
|
export_group.add_argument(
|
||
|
|
"--postprocessed-outputs",
|
||
|
|
action="store_true",
|
||
|
|
help="Export postprocessed selected outputs instead of raw head tensors.",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--include-anchor-stride",
|
||
|
|
action="store_true",
|
||
|
|
help="Also export anchors and strides. Only used with --postprocessed-outputs.",
|
||
|
|
)
|
||
|
|
parser.add_argument("--skip-torchscript", action="store_true", help="Skip TorchScript export")
|
||
|
|
parser.add_argument("--skip-onnx", action="store_true", help="Skip ONNX export")
|
||
|
|
parser.add_argument("--no-fuse", action="store_true", help="Disable Conv-BN fusion before export")
|
||
|
|
return parser.parse_args()
|
||
|
|
|
||
|
|
|
||
|
|
def _make_example_input(imgsz_wh: tuple[int, int]) -> torch.Tensor:
|
||
|
|
width, height = imgsz_wh
|
||
|
|
return torch.zeros(1, 3, height, width, dtype=torch.float32)
|
||
|
|
|
||
|
|
|
||
|
|
def _prepare_model_for_export(model: nn.Module, max_det: int, dynamic: bool, fuse: bool) -> nn.Module:
|
||
|
|
model = deepcopy(model).cpu()
|
||
|
|
for parameter in model.parameters():
|
||
|
|
parameter.requires_grad_(False)
|
||
|
|
model.eval()
|
||
|
|
model.float()
|
||
|
|
if fuse and hasattr(model, "fuse"):
|
||
|
|
model = model.fuse()
|
||
|
|
|
||
|
|
for module in model.modules():
|
||
|
|
if isinstance(module, (Detect, RTDETRDecoder)):
|
||
|
|
if hasattr(module, "dynamic"):
|
||
|
|
module.dynamic = dynamic
|
||
|
|
if hasattr(module, "max_det"):
|
||
|
|
module.max_det = int(max_det)
|
||
|
|
if hasattr(module, "shape"):
|
||
|
|
module.shape = None
|
||
|
|
return model
|
||
|
|
|
||
|
|
|
||
|
|
def _export_mode(args: argparse.Namespace) -> str:
|
||
|
|
if args.postprocessed_outputs:
|
||
|
|
return "postprocessed_outputs"
|
||
|
|
if args.detections_only:
|
||
|
|
return "detections_only"
|
||
|
|
if args.hybrid_outputs:
|
||
|
|
return "hybrid_outputs"
|
||
|
|
if args.denorm_branch_outputs:
|
||
|
|
return "denorm_branch_outputs"
|
||
|
|
return "raw_head_outputs"
|
||
|
|
|
||
|
|
|
||
|
|
def _forward_model_to_head_inputs(model: nn.Module, x: torch.Tensor) -> tuple[nn.Module, Any]:
|
||
|
|
y = []
|
||
|
|
for module in model.model[:-1]:
|
||
|
|
if module.f != -1:
|
||
|
|
x = y[module.f] if isinstance(module.f, int) else [x if j == -1 else y[j] for j in module.f]
|
||
|
|
x = module(x)
|
||
|
|
y.append(x if module.i in model.save else None)
|
||
|
|
|
||
|
|
head = model.model[-1]
|
||
|
|
if head.f != -1:
|
||
|
|
x = y[head.f] if isinstance(head.f, int) else [x if j == -1 else y[j] for j in head.f]
|
||
|
|
return head, x
|
||
|
|
|
||
|
|
|
||
|
|
def _flatten_detect3d_outputs(outputs: Any, include_anchor_stride: bool = False) -> tuple[torch.Tensor, ...]:
|
||
|
|
if not isinstance(outputs, (tuple, list)) or len(outputs) < 2:
|
||
|
|
raise RuntimeError("Model forward output must be `(detections, raw_preds)` for Detect3D models.")
|
||
|
|
|
||
|
|
detections, raw_preds = outputs[0], outputs[1]
|
||
|
|
if not isinstance(raw_preds, dict):
|
||
|
|
raise RuntimeError("Raw prediction payload is missing.")
|
||
|
|
|
||
|
|
one2one = raw_preds.get("one2one", raw_preds)
|
||
|
|
if not isinstance(one2one, dict):
|
||
|
|
raise RuntimeError("one2one prediction payload is missing.")
|
||
|
|
|
||
|
|
required_keys = ["preds_3d_selected", "preds_edge_selected"]
|
||
|
|
if include_anchor_stride:
|
||
|
|
required_keys += ["anchors_selected", "strides_selected"]
|
||
|
|
missing = [key for key in required_keys if one2one.get(key) is None]
|
||
|
|
if missing:
|
||
|
|
available = sorted(one2one.keys())
|
||
|
|
raise RuntimeError(f"Missing Detect3D export tensors {missing}. Available keys: {available}.")
|
||
|
|
|
||
|
|
result = (detections, one2one["preds_3d_selected"], one2one["preds_edge_selected"])
|
||
|
|
if include_anchor_stride:
|
||
|
|
result = result + (one2one["anchors_selected"], one2one["strides_selected"])
|
||
|
|
return result
|
||
|
|
|
||
|
|
|
||
|
|
def _denorm_detect3d_branch_outputs(model: nn.Module, x: torch.Tensor) -> tuple[torch.Tensor, ...]:
|
||
|
|
head, head_inputs = _forward_model_to_head_inputs(model, x)
|
||
|
|
branch_inputs = [tensor.detach() for tensor in head_inputs] if getattr(head, "end2end", False) else head_inputs
|
||
|
|
branch = head.one2one if getattr(head, "end2end", False) else head.one2many
|
||
|
|
raw_preds = head.forward_head(branch_inputs, **branch)
|
||
|
|
if not isinstance(raw_preds, dict):
|
||
|
|
raise RuntimeError("Raw branch payload is missing.")
|
||
|
|
|
||
|
|
required_keys = ("boxes", "scores", "preds_3d", "preds_edge")
|
||
|
|
missing = [key for key in required_keys if raw_preds.get(key) is None]
|
||
|
|
if missing:
|
||
|
|
available = sorted(raw_preds.keys())
|
||
|
|
raise RuntimeError(f"Raw branch tensors {missing} are missing. Available keys: {available}.")
|
||
|
|
|
||
|
|
return raw_preds["boxes"], raw_preds["scores"], raw_preds["preds_3d"], raw_preds["preds_edge"]
|
||
|
|
|
||
|
|
|
||
|
|
def _collect_raw_detect3d_head_outputs(head: nn.Module, head_inputs: list[torch.Tensor]) -> tuple[torch.Tensor, ...]:
|
||
|
|
if not getattr(head, "end2end", False):
|
||
|
|
raise RuntimeError("Raw-head export currently expects an end2end Detect3D head.")
|
||
|
|
|
||
|
|
branch = head.one2one
|
||
|
|
box_head = branch.get("box_head")
|
||
|
|
cls_head = branch.get("cls_head")
|
||
|
|
head_3d = branch.get("head_3d")
|
||
|
|
edge_head = branch.get("edge_head")
|
||
|
|
if any(module is None for module in (box_head, cls_head, head_3d, edge_head)):
|
||
|
|
raise RuntimeError("Raw-head export could not find complete one2one head modules.")
|
||
|
|
|
||
|
|
bs = head_inputs[0].shape[0]
|
||
|
|
raw_boxes = torch.cat([box_head[i](head_inputs[i]).view(bs, 4 * head.reg_max, -1) for i in range(head.nl)], dim=-1)
|
||
|
|
raw_scores = torch.cat([cls_head[i](head_inputs[i]).view(bs, head.nc, -1) for i in range(head.nl)], dim=-1)
|
||
|
|
raw_3d = torch.cat([head_3d[i](head_inputs[i]).view(bs, head.no_3d, -1) for i in range(head.nl)], dim=-1)
|
||
|
|
raw_edge = torch.cat([edge_head[i](head_inputs[i]).view(bs, head.no_edge, -1) for i in range(head.nl)], dim=-1)
|
||
|
|
return raw_boxes, raw_scores, raw_3d, raw_edge
|
||
|
|
|
||
|
|
|
||
|
|
def _raw_detect3d_head_outputs(model: nn.Module, x: torch.Tensor) -> tuple[torch.Tensor, ...]:
|
||
|
|
head, head_inputs = _forward_model_to_head_inputs(model, x)
|
||
|
|
return _collect_raw_detect3d_head_outputs(head, head_inputs)
|
||
|
|
|
||
|
|
|
||
|
|
def _hybrid_detect2d_raw3d_outputs(model: nn.Module, x: torch.Tensor) -> tuple[torch.Tensor, ...]:
|
||
|
|
head, head_inputs = _forward_model_to_head_inputs(model, x)
|
||
|
|
raw_boxes, raw_scores, raw_3d, raw_edge = _collect_raw_detect3d_head_outputs(head, head_inputs)
|
||
|
|
preds = {"boxes": raw_boxes, "scores": raw_scores, "feats": head_inputs}
|
||
|
|
detections = head.postprocess(head._inference(preds).permute(0, 2, 1))
|
||
|
|
return detections, raw_boxes, raw_scores, raw_3d, raw_edge
|
||
|
|
|
||
|
|
|
||
|
|
def _detections_only_outputs(model: nn.Module, x: torch.Tensor) -> tuple[torch.Tensor]:
|
||
|
|
outputs = model(x)
|
||
|
|
if not isinstance(outputs, (tuple, list)) or len(outputs) < 1:
|
||
|
|
raise RuntimeError("Model forward output must contain detections for detections-only export.")
|
||
|
|
return (outputs[0],)
|
||
|
|
|
||
|
|
|
||
|
|
class SingleROIExportWrapper(nn.Module):
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
model: nn.Module,
|
||
|
|
include_anchor_stride: bool = False,
|
||
|
|
detections_only: bool = False,
|
||
|
|
hybrid_outputs: bool = False,
|
||
|
|
denorm_branch_outputs: bool = False,
|
||
|
|
postprocessed_outputs: bool = False,
|
||
|
|
):
|
||
|
|
super().__init__()
|
||
|
|
self.model = model
|
||
|
|
self.include_anchor_stride = include_anchor_stride
|
||
|
|
self.detections_only = detections_only
|
||
|
|
self.hybrid_outputs = hybrid_outputs
|
||
|
|
self.denorm_branch_outputs = denorm_branch_outputs
|
||
|
|
self.postprocessed_outputs = postprocessed_outputs
|
||
|
|
|
||
|
|
def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, ...]:
|
||
|
|
if self.postprocessed_outputs:
|
||
|
|
return _flatten_detect3d_outputs(self.model(x), include_anchor_stride=self.include_anchor_stride)
|
||
|
|
if self.detections_only:
|
||
|
|
return _detections_only_outputs(self.model, x)
|
||
|
|
if self.hybrid_outputs:
|
||
|
|
return _hybrid_detect2d_raw3d_outputs(self.model, x)
|
||
|
|
if self.denorm_branch_outputs:
|
||
|
|
return _denorm_detect3d_branch_outputs(self.model, x)
|
||
|
|
return _raw_detect3d_head_outputs(self.model, x)
|
||
|
|
|
||
|
|
|
||
|
|
def _output_names(
|
||
|
|
include_anchor_stride: bool = False,
|
||
|
|
detections_only: bool = False,
|
||
|
|
hybrid_outputs: bool = False,
|
||
|
|
denorm_branch_outputs: bool = False,
|
||
|
|
postprocessed_outputs: bool = False,
|
||
|
|
) -> list[str]:
|
||
|
|
if postprocessed_outputs:
|
||
|
|
names = ["detections", "preds_3d", "preds_edge"]
|
||
|
|
if include_anchor_stride:
|
||
|
|
names += ["anchors", "strides"]
|
||
|
|
return names
|
||
|
|
if detections_only:
|
||
|
|
return ["detections"]
|
||
|
|
if hybrid_outputs:
|
||
|
|
return ["detections", "boxes_head_raw", "scores_head_raw", "preds_3d_head_raw", "preds_edge_head_raw"]
|
||
|
|
if denorm_branch_outputs:
|
||
|
|
return ["boxes_raw", "scores_raw", "preds_3d_raw", "preds_edge_raw"]
|
||
|
|
return ["boxes_head_raw", "scores_head_raw", "preds_3d_head_raw", "preds_edge_head_raw"]
|
||
|
|
|
||
|
|
|
||
|
|
def _describe_outputs(outputs: tuple[torch.Tensor, ...]) -> list[list[int]]:
|
||
|
|
return [list(output.shape) for output in outputs]
|
||
|
|
|
||
|
|
|
||
|
|
def _dynamic_axes(input_names: list[str], output_names: list[str], enabled: bool) -> dict[str, dict[int, str]] | None:
|
||
|
|
if not enabled:
|
||
|
|
return None
|
||
|
|
dynamic_axes = {name: {0: "batch"} for name in input_names}
|
||
|
|
dynamic_axes.update({name: {0: "batch"} for name in output_names})
|
||
|
|
return dynamic_axes
|
||
|
|
|
||
|
|
|
||
|
|
def _write_manifest(path: Path, payload: dict[str, Any]) -> None:
|
||
|
|
path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
|
||
|
|
|
||
|
|
|
||
|
|
def _base_manifest(args: argparse.Namespace, output_names: list[str], output_shapes: list[list[int]]) -> dict[str, Any]:
|
||
|
|
return {
|
||
|
|
"artifact_name": "single_roi_model",
|
||
|
|
"model_path": str(Path(args.model_path).resolve()),
|
||
|
|
"input_names": ["images"],
|
||
|
|
"input_sizes_wh": {"images": [int(args.imgsz[0]), int(args.imgsz[1])]},
|
||
|
|
"output_names": output_names,
|
||
|
|
"output_shapes": output_shapes,
|
||
|
|
"dynamic_batch": bool(args.dynamic),
|
||
|
|
"max_det": int(args.max_det),
|
||
|
|
"opset": int(args.opset),
|
||
|
|
"fuse": not bool(args.no_fuse),
|
||
|
|
"include_anchor_stride": bool(args.include_anchor_stride),
|
||
|
|
"detections_only": bool(args.detections_only),
|
||
|
|
"hybrid_outputs": bool(args.hybrid_outputs),
|
||
|
|
"denorm_branch_outputs": bool(args.denorm_branch_outputs),
|
||
|
|
"postprocessed_outputs": bool(args.postprocessed_outputs),
|
||
|
|
"export_mode": _export_mode(args),
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def _export_torchscript(model: nn.Module, example_input: torch.Tensor, save_path: Path, manifest: dict[str, Any]) -> None:
|
||
|
|
LOGGER.info(f"Exporting TorchScript to {save_path}")
|
||
|
|
traced = torch.jit.trace(model, example_input, strict=False)
|
||
|
|
traced.save(str(save_path), _extra_files={"config.txt": json.dumps(manifest, sort_keys=True)})
|
||
|
|
LOGGER.info(f"TorchScript export success: {save_path} ({save_path.stat().st_size / (1024 * 1024):.2f} MB)")
|
||
|
|
|
||
|
|
|
||
|
|
def _save_onnx_metadata(onnx_path: Path, manifest: dict[str, Any], simplify: bool) -> None:
|
||
|
|
import onnx
|
||
|
|
|
||
|
|
model_onnx = onnx.load(str(onnx_path))
|
||
|
|
for key, value in manifest.items():
|
||
|
|
meta = model_onnx.metadata_props.add()
|
||
|
|
meta.key, meta.value = key, json.dumps(value, ensure_ascii=True) if isinstance(value, (dict, list)) else str(value)
|
||
|
|
if simplify:
|
||
|
|
import onnxslim
|
||
|
|
|
||
|
|
LOGGER.info("Simplifying ONNX graph with onnxslim")
|
||
|
|
model_onnx = onnxslim.slim(model_onnx)
|
||
|
|
onnx.save(model_onnx, str(onnx_path))
|
||
|
|
|
||
|
|
|
||
|
|
def _export_onnx(
|
||
|
|
model: nn.Module,
|
||
|
|
example_input: torch.Tensor,
|
||
|
|
save_path: Path,
|
||
|
|
output_names: list[str],
|
||
|
|
manifest: dict[str, Any],
|
||
|
|
opset: int,
|
||
|
|
dynamic: bool,
|
||
|
|
simplify: bool,
|
||
|
|
) -> None:
|
||
|
|
LOGGER.info(f"Exporting ONNX to {save_path} with opset={opset}")
|
||
|
|
patch_args = SimpleNamespace(dynamic=dynamic, half=False, format="onnx")
|
||
|
|
with onnx_export_patch():
|
||
|
|
with arange_patch(patch_args):
|
||
|
|
torch.onnx.export(
|
||
|
|
model,
|
||
|
|
example_input,
|
||
|
|
str(save_path),
|
||
|
|
verbose=False,
|
||
|
|
opset_version=opset,
|
||
|
|
do_constant_folding=True,
|
||
|
|
input_names=["images"],
|
||
|
|
output_names=output_names,
|
||
|
|
dynamic_axes=_dynamic_axes(["images"], output_names, dynamic),
|
||
|
|
)
|
||
|
|
try:
|
||
|
|
_save_onnx_metadata(save_path, manifest, simplify=simplify)
|
||
|
|
except ImportError as error:
|
||
|
|
LOGGER.warning(f"Skipping ONNX metadata/simplify step because a package is missing: {error}")
|
||
|
|
except Exception as error:
|
||
|
|
LOGGER.warning(f"ONNX post-processing failed: {error}")
|
||
|
|
LOGGER.info(f"ONNX export success: {save_path} ({save_path.stat().st_size / (1024 * 1024):.2f} MB)")
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> None:
|
||
|
|
args = parse_args()
|
||
|
|
if args.include_anchor_stride and not args.postprocessed_outputs:
|
||
|
|
LOGGER.warning("--include-anchor-stride is ignored unless --postprocessed-outputs is also set.")
|
||
|
|
LOGGER.info(f"Export mode: {_export_mode(args)}")
|
||
|
|
|
||
|
|
save_dir = Path(args.save_dir)
|
||
|
|
save_dir.mkdir(parents=True, exist_ok=True)
|
||
|
|
|
||
|
|
model, _ = load_checkpoint(args.model_path, device="cpu", fuse=False)
|
||
|
|
prepared_model = _prepare_model_for_export(model, max_det=args.max_det, dynamic=args.dynamic, fuse=not args.no_fuse)
|
||
|
|
export_model = SingleROIExportWrapper(
|
||
|
|
prepared_model,
|
||
|
|
include_anchor_stride=args.include_anchor_stride,
|
||
|
|
detections_only=args.detections_only,
|
||
|
|
hybrid_outputs=args.hybrid_outputs,
|
||
|
|
denorm_branch_outputs=args.denorm_branch_outputs,
|
||
|
|
postprocessed_outputs=args.postprocessed_outputs,
|
||
|
|
).eval()
|
||
|
|
|
||
|
|
example_input = _make_example_input((int(args.imgsz[0]), int(args.imgsz[1])))
|
||
|
|
with torch.inference_mode():
|
||
|
|
outputs = export_model(example_input)
|
||
|
|
|
||
|
|
output_names = _output_names(
|
||
|
|
include_anchor_stride=args.include_anchor_stride,
|
||
|
|
detections_only=args.detections_only,
|
||
|
|
hybrid_outputs=args.hybrid_outputs,
|
||
|
|
denorm_branch_outputs=args.denorm_branch_outputs,
|
||
|
|
postprocessed_outputs=args.postprocessed_outputs,
|
||
|
|
)
|
||
|
|
LOGGER.info(f"Single ROI model output shapes: {_describe_outputs(outputs)}")
|
||
|
|
|
||
|
|
manifest = _base_manifest(args=args, output_names=output_names, output_shapes=_describe_outputs(outputs))
|
||
|
|
_write_manifest(save_dir / "single_roi_model.export.json", manifest)
|
||
|
|
|
||
|
|
if not args.skip_torchscript:
|
||
|
|
_export_torchscript(export_model, example_input, save_dir / "single_roi_model.torchscript", manifest)
|
||
|
|
if not args.skip_onnx:
|
||
|
|
_export_onnx(
|
||
|
|
export_model,
|
||
|
|
example_input,
|
||
|
|
save_dir / "single_roi_model.onnx",
|
||
|
|
output_names=output_names,
|
||
|
|
manifest=manifest,
|
||
|
|
opset=args.opset,
|
||
|
|
dynamic=args.dynamic,
|
||
|
|
simplify=args.simplify,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|