单目3D初始代码

tools/model_merging/export_single_roi_yolo26.py

@@ -0,0 +1,408 @@
+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()

tools/model_merging/export_single_roi_yolo26.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
+export PYTHONPATH="${PROJECT_ROOT}:${PYTHONPATH}"
+
+python tools/model_merging/export_single_roi_yolo26.py \
+    --model-path runs/detect/train_mono3d_roi0_202603291330_epoch61.pt \
+    --save-dir runs/export/train_mono3d_single_roi0_202603291330-postprocessed-original \
+    --imgsz 768 352 \
+    --opset 11 \
+    --postprocessed-outputs
+    # --detections-only
+
+# Example: export ROI1 instead
+# python tools/model_merging/export_single_roi_yolo26.py \
+#     --model-path runs/detect/train_mono3d_roi1_202603291330_epoch99.pt \
+#     --save-dir runs/export/train_mono3d_single_roi1_202603291330-detections \
+#     --imgsz 768 352 \
+#     --opset 15 \
+#     --detections-only

tools/model_merging/merge_models_of_2roi.py

@@ -0,0 +1,329 @@
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+FILE = Path(__file__).resolve()
+ROOT = FILE.parents[1]
+if str(ROOT) not in sys.path:
+    sys.path.append(str(ROOT))
+
+from models.common import DetectMultiBackend
+from export import attempt_load
+from utils.general import LOGGER
+
+
+def analyze_checkpoint(weights_path):
+    """Analyze checkpoint file structure to understand size."""
+    import torch
+    
+    ckpt = torch.load(weights_path, map_location='cpu')
+    
+    LOGGER.info(f"\n{'='*80}")
+    LOGGER.info(f"CHECKPOINT ANALYSIS: {Path(weights_path).name}")
+    LOGGER.info(f"{'='*80}")
+    
+    # Check what's in the checkpoint
+    if isinstance(ckpt, dict):
+        LOGGER.info(f"Checkpoint keys: {list(ckpt.keys())}")
+        
+        # Analyze each component size
+        total_params = 0
+        for key in ckpt.keys():
+            if key in ['model', 'ema']:
+                model_obj = ckpt[key]
+                if hasattr(model_obj, 'parameters'):
+                    params = sum(p.numel() for p in model_obj.parameters())
+                    param_size_mb = params * 4 / (1024**2)  # FP32 = 4 bytes
+                    total_params += params
+                    LOGGER.info(f"  - {key}: {params:,} parameters ({param_size_mb:.2f}MB if FP32)")
+        
+        LOGGER.info(f"Total parameters: {total_params:,}")
+        LOGGER.info(f"Theoretical FP32 size: {total_params * 4 / (1024**2):.2f}MB")
+        
+        # Check actual file size
+        import os
+        actual_size = os.path.getsize(weights_path) / (1024**2)
+        compression_ratio = (total_params * 4 / (1024**2)) / actual_size if actual_size > 0 else 0
+        LOGGER.info(f"Actual file size: {actual_size:.2f}MB")
+        LOGGER.info(f"Compression ratio: {compression_ratio:.2f}x")
+        
+    LOGGER.info(f"{'='*80}\n")
+
+
+def load_model_lightweight(weights_path, device='cpu'):
+    """Load model in the most lightweight way possible (state_dict only).
+    
+    This avoids loading EMA weights and other training artifacts that can double the size.
+    """
+    import torch
+    from models.yolo import Model
+    
+    LOGGER.info(f"Loading model with lightweight method: {weights_path}")
+    ckpt = torch.load(weights_path, map_location=device)
+    
+    # Prefer regular model over EMA (EMA weights can be larger)
+    if 'model' in ckpt:
+        model = ckpt['model']
+    elif 'ema' in ckpt:
+        LOGGER.warning("Only EMA weights available, using them (may increase size)")
+        model = ckpt['ema']
+    else:
+        raise ValueError(f"No model found in checkpoint: {weights_path}")
+    
+    # Convert to float32 and eval mode
+    model = model.float().eval()
+    
+    # Remove training-specific attributes to reduce memory
+    if hasattr(model, 'hyp'):
+        delattr(model, 'hyp')
+    if hasattr(model, 'gr'):
+        delattr(model, 'gr')
+    
+    return model
+
+
+class Model_Merged(nn.Module):
+    def __init__(self, model_roi0, model_roi1):
+        super(Model_Merged, self).__init__()
+        self.model_roi0 = model_roi0
+        self.model_roi1 = model_roi1
+
+        self.model_roi0.model[-1].export_raw = True  # set export=True for Detect head
+        self.model_roi1.model[-1].export_raw = True  # set export=True for Detect
+
+        # 确保两个模型都设置为eval模式
+        self.model_roi0.eval()
+        self.model_roi1.eval()
+
+        self.stride = model_roi0.stride  # assuming both models have the same stride
+        self.names = model_roi0.names  # assuming both models have the same class names
+
+    def forward(self, x_roi0, x_roi1):
+
+        with torch.no_grad():
+            out_roi0 = self.model_roi0(x_roi0)
+            out_roi1 = self.model_roi1(x_roi1)
+
+        return out_roi0 + out_roi1
+
+def merge_models(roi0_model_path, roi1_model_path, save_path, use_lightweight=False):
+    """Load and merge two models, ensuring minimal memory footprint.
+    
+    Issues with large model size:
+    1. attempt_load may load EMA weights which are larger
+    2. fuse=True can sometimes increase size due to optimization overhead
+    3. The checkpoint contains training metadata
+    
+    Args:
+        use_lightweight: If True, use custom lightweight loading (experimental)
+    """
+    import os
+    
+    # Report original model sizes
+    size_roi0 = os.path.getsize(roi0_model_path) / (1024 * 1024)  # MB
+    size_roi1 = os.path.getsize(roi1_model_path) / (1024 * 1024)  # MB
+    LOGGER.info(f"Original model sizes: ROI0={size_roi0:.2f}MB, ROI1={size_roi1:.2f}MB, Total={size_roi0+size_roi1:.2f}MB")
+    
+    # Load the two models - try to minimize size
+    if use_lightweight:
+        LOGGER.info("Using lightweight model loading (experimental)")
+        model_roi0 = load_model_lightweight(roi0_model_path, device='cpu')
+        model_roi1 = load_model_lightweight(roi1_model_path, device='cpu')
+    else:
+        LOGGER.info(f"Loading ROI0 model: {roi0_model_path}")
+        model_roi0 = attempt_load(roi0_model_path, device='cpu', fuse=False)  # fuse=False to avoid overhead
+        
+        LOGGER.info(f"Loading ROI1 model: {roi1_model_path}")
+        model_roi1 = attempt_load(roi1_model_path, device='cpu', fuse=False)
+
+    # Create merged model
+    merged_model = Model_Merged(model_roi0, model_roi1)
+
+    return merged_model
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--roi0_model_path', type=str, default='release/yolov5s-30w/roi0_last.pt', help='Path to the ROI0 model file')
+    parser.add_argument('--roi1_model_path', type=str, default='release/yolov5s-30w/roi1_last.pt', help='Path to the ROI1 model file')
+    parser.add_argument('--save_dir', type=str, default='release/yolov5s-30w', help='Path to save the merged model file')
+    parser.add_argument('--imgsz', nargs=2, type=int, default=[352, 704], metavar=('H', 'W'), help='Input image size (height width), e.g. --imgsz 352 704')
+    parser.add_argument('--opset', type=int, default=11, help='ONNX opset version (11, 12, 13, 17, etc.)')
+    parser.add_argument('--skip-export', action='store_true', help='Skip exporting to TorchScript/ONNX')
+    parser.add_argument('--lightweight', action='store_true', help='Use lightweight model loading (may reduce export size)')
+    parser.add_argument('--export-separate', action='store_true', help='Export two models separately instead of merged (saves space)')
+    parser.add_argument('--analyze', action='store_true', help='Analyze checkpoint structure and exit')
+    args = parser.parse_args()
+    
+    # Analyze checkpoint if requested
+    if args.analyze:
+        analyze_checkpoint(args.roi0_model_path)
+        if args.roi0_model_path != args.roi1_model_path:
+            analyze_checkpoint(args.roi1_model_path)
+        sys.exit(0)
+
+    if not Path(args.save_dir).exists():
+        Path(args.save_dir).mkdir(parents=True, exist_ok=True)
+
+    merged_model = merge_models(args.roi0_model_path, args.roi1_model_path, args.save_dir, args.lightweight)
+    merged_model.eval()
+
+    h, w = args.imgsz
+    im0 = torch.randn(1, 3, h, w)  # example input for roi0
+    im1 = torch.randn(1, 3, h, w)
+
+    output = merged_model(im0, im1)
+
+    print(f"Merged model created successfully.")
+    
+    # Important note about model size
+    LOGGER.info("=" * 80)
+    LOGGER.info("MODEL SIZE EXPLANATION:")
+    LOGGER.info("The merged model contains TWO complete and independent models.")
+    LOGGER.info("Expected export size: ~2x original size (≈100MB for 2x25MB models)")
+    LOGGER.info("This is NORMAL because the two models have different trained weights.")
+    if args.export_separate:
+        LOGGER.info("Using --export-separate to save space by exporting separately.")
+    LOGGER.info("=" * 80)
+    
+    if args.skip_export:
+        LOGGER.info("Skipping export (--skip-export flag set)")
+        sys.exit(0)
+    
+    # Option 1: Export separately (recommended for saving space)
+    if args.export_separate:
+        LOGGER.info("Exporting models SEPARATELY to save space...")
+        
+        # Export ROI0 model
+        save_path_roi0_ts = Path(args.save_dir) / "roi0_model.torchscript"
+        save_path_roi0_onnx = Path(args.save_dir) / "roi0_model.onnx"
+        
+        try:
+            LOGGER.info(f"Exporting ROI0 to TorchScript: {save_path_roi0_ts}")
+            ts0 = torch.jit.trace(merged_model.model_roi0, im0, strict=False)
+            ts0.save(str(save_path_roi0_ts))
+            size_ts0 = save_path_roi0_ts.stat().st_size / (1024 * 1024)
+            LOGGER.info(f"ROI0 TorchScript export success: {size_ts0:.2f}MB")
+            
+            LOGGER.info(f"Exporting ROI0 to ONNX: {save_path_roi0_onnx}")
+            torch.onnx.export(
+                merged_model.model_roi0, im0, str(save_path_roi0_onnx),
+                verbose=False, opset_version=args.opset, do_constant_folding=True,
+                input_names=["input"], output_names=["output"]
+            )
+            size_onnx0 = save_path_roi0_onnx.stat().st_size / (1024 * 1024)
+            LOGGER.info(f"ROI0 ONNX export success: {size_onnx0:.2f}MB")
+        except Exception as e:
+            LOGGER.error(f"ROI0 export failed: {e}")
+        
+        # Export ROI1 model
+        save_path_roi1_ts = Path(args.save_dir) / "roi1_model.torchscript"
+        save_path_roi1_onnx = Path(args.save_dir) / "roi1_model.onnx"
+        
+        try:
+            LOGGER.info(f"Exporting ROI1 to TorchScript: {save_path_roi1_ts}")
+            ts1 = torch.jit.trace(merged_model.model_roi1, im1, strict=False)
+            ts1.save(str(save_path_roi1_ts))
+            size_ts1 = save_path_roi1_ts.stat().st_size / (1024 * 1024)
+            LOGGER.info(f"ROI1 TorchScript export success: {size_ts1:.2f}MB")
+            
+            LOGGER.info(f"Exporting ROI1 to ONNX: {save_path_roi1_onnx}")
+            torch.onnx.export(
+                merged_model.model_roi1, im1, str(save_path_roi1_onnx),
+                verbose=False, opset_version=args.opset, do_constant_folding=True,
+                input_names=["input"], output_names=["output"]
+            )
+            size_onnx1 = save_path_roi1_onnx.stat().st_size / (1024 * 1024)
+            LOGGER.info(f"ROI1 ONNX export success: {size_onnx1:.2f}MB")
+        except Exception as e:
+            LOGGER.error(f"ROI1 export failed: {e}")
+        
+        LOGGER.info("=" * 80)
+        LOGGER.info("SEPARATE EXPORT COMPLETE")
+        LOGGER.info("Use these files for inference by loading both models separately")
+        LOGGER.info("=" * 80)
+        sys.exit(0)
+
+    # Option 2: Export merged model (default, but larger)
+    save_path_ts = Path(args.save_dir) / "merged_model.torchscript"
+    try:
+        LOGGER.info(f"Exporting to TorchScript: {save_path_ts}")
+        # Use torch.jit.trace with tuple of inputs for multi-input models
+        ts = torch.jit.trace(merged_model, (im0, im1), strict=False)
+        ts.save(str(save_path_ts))
+        
+        # Report file size
+        size_ts = save_path_ts.stat().st_size / (1024 * 1024)  # MB
+        LOGGER.info(f"TorchScript export success: {save_path_ts} (Size: {size_ts:.2f}MB)")
+    except Exception as e:
+        LOGGER.error(f"TorchScript export failed: {e}")
+
+    # Export to ONNX
+    save_path_onnx = Path(args.save_dir) / "merged_model.onnx"
+    try:
+        # Check PyTorch version to determine compatible opset version
+        import torch
+        torch_version = torch.__version__
+        LOGGER.info(f"PyTorch version: {torch_version}")
+        
+        # Use opset 11 or 13 for better compatibility, or accept the default (18)
+        # opset 11: widely supported by TensorRT, OpenVINO, ONNX Runtime
+        # opset 13: better operator coverage
+        # opset 17+: required for newer PyTorch versions
+        opset = args.opset  # Use command line argument
+        
+        LOGGER.info(f"Exporting to ONNX with opset_version={opset}: {save_path_onnx}")
+        torch.onnx.export(
+            merged_model,
+            (im0, im1),  # tuple of example inputs
+            str(save_path_onnx),
+            verbose=False,
+            opset_version=opset,
+            do_constant_folding=True,
+            input_names=["roi0_input", "roi1_input"],  # specify names for two inputs
+            output_names=["roi0_output_2d_8x", "roi0_output_2d_16x", "roi0_output_2d_32x", 
+                          "roi0_output_3d_8x", "roi0_output_3d_16x", "roi0_output_3d_32x",
+                          "roi1_output_2d_8x", "roi1_output_2d_16x", "roi1_output_2d_32x", 
+                          "roi1_output_3d_8x", "roi1_output_3d_16x", "roi1_output_3d_32x"],  # specify names for two outputs
+            dynamic_axes={
+                "roi0_input": {0: "batch"},
+                "roi1_input": {0: "batch"},
+                "roi0_output": {0: "batch"},
+                "roi1_output": {0: "batch"}
+            } if False else None  # set to True to enable dynamic batch size
+        )
+        LOGGER.info(f"ONNX export success: {save_path_onnx}")
+        
+        # Report file size before simplification
+        size_onnx_before = save_path_onnx.stat().st_size / (1024 * 1024)  # MB
+        LOGGER.info(f"ONNX model size before simplification: {size_onnx_before:.2f}MB")
+        
+        # Optional: simplify ONNX model
+        try:
+            import onnx
+            import onnxslim
+            
+            # Skip simplification for multi-input models as it may increase size
+            LOGGER.warning("Skipping ONNX simplification for multi-input model (often increases size)")
+            LOGGER.info("If you still want to try simplification, use: onnxslim input.onnx output.onnx")
+            
+            # Uncomment below to force simplification (not recommended for multi-input models)
+            # LOGGER.info("Simplifying ONNX model with onnxslim...")
+            # model_onnx = onnx.load(str(save_path_onnx))
+            # model_onnx = onnxslim.slim(model_onnx)
+            # onnx.save(model_onnx, str(save_path_onnx))
+            # 
+            # # Report size after simplification
+            # size_onnx_after = save_path_onnx.stat().st_size / (1024 * 1024)  # MB
+            # reduction = size_onnx_before - size_onnx_after
+            # LOGGER.info(f"ONNX simplification success: {size_onnx_after:.2f}MB (reduced by {reduction:.2f}MB)")
+        except ImportError:
+            LOGGER.warning("onnxslim not installed, skipping simplification. Install with: pip install onnxslim")
+        except Exception as e:
+            LOGGER.warning(f"ONNX simplification failed: {e}")
+    except Exception as e:
+        LOGGER.error(f"ONNX export failed: {e}")

tools/model_merging/merge_models_of_2roi.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+# 
+# Merge two ROI models into a single model
+#
+# Arguments:
+#   --roi0_model_path: Path to the ROI0 model file
+#   --roi1_model_path: Path to the ROI1 model file
+#   --save_dir: Directory to save the merged model file
+
+
+
+# Set PYTHONPATH to project root for module imports
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
+export PYTHONPATH="${PROJECT_ROOT}:${PYTHONPATH}"
+python tools/model_merging/merge_models_of_2roi.py \
+    --roi0_model_path runs/train/yolov5s-300w-roi0-newdata-cncap-bottom_aug_20260312_012251/weights/last.pt \
+    --roi1_model_path runs/train/yolov5s-300w-roi1-newdata-cncap-bottom_aug_20260312_012755/weights/last.pt \
+    --save_dir release/yolov5s-300w-newdata-cncap-768x352 \
+    --imgsz 352 768

tools/model_merging/merge_models_of_2roi_yolo26.py

@@ -0,0 +1,740 @@
+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_ROI0_MODEL = ROOT / "runs" / "detect" / "mono3d_roi0_20260506_epoch99.pt"
+DEFAULT_ROI1_MODEL = ROOT / "runs" / "detect" / "mono3d_roi1_20260506_epoch99.pt"
+DEFAULT_SAVE_DIR = ROOT / "runs" / "export" / "train_mono3d_two_roi_20260506"
+DEFAULT_IMGSZ_WH = (768, 352)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Merge two yolo26 single-ROI Detect3D checkpoints and export them.")
+    parser.add_argument("--roi0-model-path", type=str, default=str(DEFAULT_ROI0_MODEL), help="Path to ROI0 checkpoint")
+    parser.add_argument("--roi1-model-path", type=str, default=str(DEFAULT_ROI1_MODEL), help="Path to ROI1 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="Common example input size as width height, e.g. --imgsz 768 352",
+    )
+    parser.add_argument("--roi0-imgsz", nargs=2, type=int, default=None, metavar=("W", "H"), help="Optional ROI0 input size override")
+    parser.add_argument("--roi1-imgsz", nargs=2, type=int, default=None, metavar=("W", "H"), help="Optional ROI1 input size override")
+    parser.add_argument("--opset", type=int, default=17, help="ONNX opset version")
+    parser.add_argument("--max-det", type=int, default=300, help="Top-k detections kept by each 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(
+        "--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 for each ROI. Only used with --postprocessed-outputs.",
+    )
+    parser.add_argument("--export-separate", action="store_true", help="Export ROI0 and ROI1 as separate single-input artifacts")
+    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")
+    parser.add_argument(
+        "--edge-head-mode",
+        type=str,
+        choices=("keep", "drop"),
+        default="keep",
+        help="Control whether edge_head branch outputs are kept in exported artifacts.",
+    )
+    parser.add_argument(
+        "--fake-3d-branch-mode",
+        type=str,
+        choices=("keep", "drop"),
+        default="keep",
+        help="Control whether fake 3D branch outputs are kept in exported artifacts.",
+    )
+    return parser.parse_args()
+
+
+def _resolve_imgsz_wh(common_imgsz_wh: tuple[int, int], override: list[int] | tuple[int, int] | None) -> tuple[int, int]:
+    if override is None:
+        return int(common_imgsz_wh[0]), int(common_imgsz_wh[1])
+    return int(override[0]), int(override[1])
+
+
+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.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,
+    roi_name: str,
+    include_anchor_stride: bool = False,
+    keep_edge_head: bool = True,
+    keep_fake_3d_branch: bool = True,
+) -> tuple[torch.Tensor, ...]:
+    if not isinstance(outputs, (tuple, list)) or len(outputs) < 2:
+        raise RuntimeError(f"{roi_name} 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(f"{roi_name} raw prediction payload is missing.")
+
+    one2one = raw_preds.get("one2one", raw_preds)
+    if not isinstance(one2one, dict):
+        raise RuntimeError(f"{roi_name} one2one prediction payload is missing.")
+
+    required_keys = ["preds_3d_selected"]
+    if keep_edge_head:
+        required_keys.append("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"{roi_name} is missing Detect3D export tensors {missing}. Available keys: {available}. "
+            "This script expects yolo26 end2end Detect3D checkpoints."
+        )
+
+    result = (
+        detections,
+        one2one["preds_3d_selected"],
+        one2one.get("preds_diff_selected"),
+    )
+    if keep_fake_3d_branch:
+        result = result + (one2one.get("preds_3d_fake_selected"),)
+    if keep_edge_head:
+        result = result + (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,
+    roi_name: str,
+    keep_edge_head: bool = True,
+    keep_fake_3d_branch: bool = True,
+) -> tuple[torch.Tensor, ...]:
+    head, head_inputs = _forward_model_to_head_inputs(model, x)
+    if not hasattr(head, "forward_head"):
+        raise RuntimeError(f"{roi_name} final layer does not expose forward_head(), cannot export raw branch outputs.")
+
+    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(f"{roi_name} raw branch payload is missing.")
+
+    required_keys = ["boxes", "scores", "preds_3d"]
+    if keep_edge_head:
+        required_keys.append("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"{roi_name} raw branch tensors {missing} are missing. Available keys: {available}.")
+
+    result = (
+        raw_preds["boxes"],
+        raw_preds["scores"],
+        raw_preds["preds_3d"],
+        raw_preds.get("preds_diff"),
+    )
+    if keep_fake_3d_branch:
+        result = result + (raw_preds.get("preds_3d_fake"),)
+    if keep_edge_head:
+        result = result + (raw_preds["preds_edge"],)
+    return result
+
+
+def _collect_raw_detect3d_head_outputs(
+    head: nn.Module,
+    head_inputs: list[torch.Tensor],
+    roi_name: str,
+    keep_edge_head: bool = True,
+    keep_fake_3d_branch: bool = True,
+) -> tuple[torch.Tensor, ...]:
+    if not getattr(head, "end2end", False):
+        raise RuntimeError(f"{roi_name} raw-head export currently expects an end2end Detect3D head.")
+
+    branch = head.one2one
+    box_head = branch.get("box_head")
+    cls_head = branch.get("cls_head")
+    diff_head = branch.get("diff_head")
+    head_3d = branch.get("head_3d")
+    fake_head_3d = branch.get("fake_head_3d")
+    edge_head = branch.get("edge_head")
+    required_modules = (
+        box_head,
+        cls_head,
+        diff_head,
+        head_3d,
+        edge_head,
+    ) if keep_edge_head else (
+        box_head,
+        cls_head,
+        diff_head,
+        head_3d,
+    )
+    if keep_fake_3d_branch:
+        required_modules = (*required_modules, fake_head_3d)
+    if any(module is None for module in required_modules):
+        raise RuntimeError(f"{roi_name} 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_diff = torch.cat([diff_head[i](head_inputs[i]).view(bs, 1, -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)
+    result = (raw_boxes, raw_scores, raw_3d, raw_diff)
+    if keep_fake_3d_branch:
+        raw_3d_fake = torch.cat([fake_head_3d[i](head_inputs[i]).view(bs, head.no_3d, -1) for i in range(head.nl)], dim=-1)
+        result = result + (raw_3d_fake,)
+    if keep_edge_head:
+        raw_edge = torch.cat([edge_head[i](head_inputs[i]).view(bs, head.no_edge, -1) for i in range(head.nl)], dim=-1)
+        result = result + (raw_edge,)
+    return result
+
+
+def _raw_detect3d_head_outputs(
+    model: nn.Module,
+    x: torch.Tensor,
+    roi_name: str,
+    keep_edge_head: bool = True,
+    keep_fake_3d_branch: bool = True,
+) -> tuple[torch.Tensor, ...]:
+    head, head_inputs = _forward_model_to_head_inputs(model, x)
+    return _collect_raw_detect3d_head_outputs(
+        head,
+        head_inputs,
+        roi_name,
+        keep_edge_head=keep_edge_head,
+        keep_fake_3d_branch=keep_fake_3d_branch,
+    )
+
+
+def _hybrid_detect2d_raw3d_outputs(
+    model: nn.Module,
+    x: torch.Tensor,
+    roi_name: str,
+    keep_edge_head: bool = True,
+    keep_fake_3d_branch: bool = True,
+) -> tuple[torch.Tensor, ...]:
+    head, head_inputs = _forward_model_to_head_inputs(model, x)
+    raw_outputs = _collect_raw_detect3d_head_outputs(
+        head,
+        head_inputs,
+        roi_name,
+        keep_edge_head=keep_edge_head,
+        keep_fake_3d_branch=keep_fake_3d_branch,
+    )
+    raw_boxes, raw_scores, *_ = raw_outputs
+    preds = {"boxes": raw_boxes, "scores": raw_scores, "feats": head_inputs}
+    detections = head.postprocess(head._inference(preds).permute(0, 2, 1))
+    return (detections, *raw_outputs)
+
+
+class ROISelectedOutputsWrapper(nn.Module):
+    def __init__(
+        self,
+        model: nn.Module,
+        roi_name: str,
+        include_anchor_stride: bool = False,
+        hybrid_outputs: bool = False,
+        denorm_branch_outputs: bool = False,
+        postprocessed_outputs: bool = False,
+        keep_edge_head: bool = True,
+        keep_fake_3d_branch: bool = True,
+    ):
+        super().__init__()
+        self.model = model
+        self.roi_name = roi_name
+        self.include_anchor_stride = include_anchor_stride
+        self.hybrid_outputs = hybrid_outputs
+        self.denorm_branch_outputs = denorm_branch_outputs
+        self.postprocessed_outputs = postprocessed_outputs
+        self.keep_edge_head = keep_edge_head
+        self.keep_fake_3d_branch = keep_fake_3d_branch
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, ...]:
+        if self.postprocessed_outputs:
+            return _flatten_detect3d_outputs(
+                self.model(x),
+                self.roi_name,
+                include_anchor_stride=self.include_anchor_stride,
+                keep_edge_head=self.keep_edge_head,
+                keep_fake_3d_branch=self.keep_fake_3d_branch,
+            )
+        if self.hybrid_outputs:
+            return _hybrid_detect2d_raw3d_outputs(
+                self.model,
+                x,
+                self.roi_name,
+                keep_edge_head=self.keep_edge_head,
+                keep_fake_3d_branch=self.keep_fake_3d_branch,
+            )
+        if self.denorm_branch_outputs:
+            return _denorm_detect3d_branch_outputs(
+                self.model,
+                x,
+                self.roi_name,
+                keep_edge_head=self.keep_edge_head,
+                keep_fake_3d_branch=self.keep_fake_3d_branch,
+            )
+        return _raw_detect3d_head_outputs(
+            self.model,
+            x,
+            self.roi_name,
+            keep_edge_head=self.keep_edge_head,
+            keep_fake_3d_branch=self.keep_fake_3d_branch,
+        )
+
+
+class TwoROIMergedExportModel(nn.Module):
+    def __init__(
+        self,
+        roi0_model: nn.Module,
+        roi1_model: nn.Module,
+        include_anchor_stride: bool = False,
+        hybrid_outputs: bool = False,
+        denorm_branch_outputs: bool = False,
+        postprocessed_outputs: bool = False,
+        keep_edge_head: bool = True,
+        keep_fake_3d_branch: bool = True,
+    ):
+        super().__init__()
+        self.roi0_model = ROISelectedOutputsWrapper(
+            roi0_model,
+            "roi0",
+            include_anchor_stride=include_anchor_stride,
+            hybrid_outputs=hybrid_outputs,
+            denorm_branch_outputs=denorm_branch_outputs,
+            postprocessed_outputs=postprocessed_outputs,
+            keep_edge_head=keep_edge_head,
+            keep_fake_3d_branch=keep_fake_3d_branch,
+        )
+        self.roi1_model = ROISelectedOutputsWrapper(
+            roi1_model,
+            "roi1",
+            include_anchor_stride=include_anchor_stride,
+            hybrid_outputs=hybrid_outputs,
+            denorm_branch_outputs=denorm_branch_outputs,
+            postprocessed_outputs=postprocessed_outputs,
+            keep_edge_head=keep_edge_head,
+            keep_fake_3d_branch=keep_fake_3d_branch,
+        )
+
+        self.stride = getattr(roi0_model, "stride", None)
+        self.names = getattr(roi0_model, "names", None)
+
+    def forward(self, x_roi0: torch.Tensor, x_roi1: torch.Tensor) -> tuple[torch.Tensor, ...]:
+        return (*self.roi0_model(x_roi0), *self.roi1_model(x_roi1))
+
+
+def _output_names(
+    prefix: str,
+    include_anchor_stride: bool = False,
+    hybrid_outputs: bool = False,
+    denorm_branch_outputs: bool = False,
+    postprocessed_outputs: bool = False,
+    keep_edge_head: bool = True,
+    keep_fake_3d_branch: bool = True,
+) -> list[str]:
+    if postprocessed_outputs:
+        names = [
+            f"{prefix}_detections",
+            f"{prefix}_preds_3d",
+            f"{prefix}_preds_diff",
+        ]
+        if keep_fake_3d_branch:
+            names.append(f"{prefix}_preds_3d_fake")
+        if keep_edge_head:
+            names.append(f"{prefix}_preds_edge")
+        if include_anchor_stride:
+            names += [f"{prefix}_anchors", f"{prefix}_strides"]
+        return names
+
+    if hybrid_outputs:
+        names = [
+            f"{prefix}_detections",
+            f"{prefix}_boxes_head_raw",
+            f"{prefix}_scores_head_raw",
+            f"{prefix}_preds_3d_head_raw",
+            f"{prefix}_preds_diff_head_raw",
+        ]
+        if keep_fake_3d_branch:
+            names.append(f"{prefix}_preds_3d_fake_head_raw")
+        if keep_edge_head:
+            names.append(f"{prefix}_preds_edge_head_raw")
+        return names
+
+    if denorm_branch_outputs:
+        names = [
+            f"{prefix}_boxes_raw",
+            f"{prefix}_scores_raw",
+            f"{prefix}_preds_3d_raw",
+            f"{prefix}_preds_diff_raw",
+        ]
+        if keep_fake_3d_branch:
+            names.append(f"{prefix}_preds_3d_fake_raw")
+        if keep_edge_head:
+            names.append(f"{prefix}_preds_edge_raw")
+        return names
+
+    names = [
+        f"{prefix}_boxes_head_raw",
+        f"{prefix}_scores_head_raw",
+        f"{prefix}_preds_3d_head_raw",
+        f"{prefix}_preds_diff_head_raw",
+    ]
+    if keep_fake_3d_branch:
+        names.append(f"{prefix}_preds_3d_fake_head_raw")
+    if keep_edge_head:
+        names.append(f"{prefix}_preds_edge_head_raw")
+    return names
+
+
+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(
+    artifact_name: str,
+    args: argparse.Namespace,
+    input_names: list[str],
+    input_sizes_wh: dict[str, list[int]],
+    output_names: list[str],
+    output_shapes: list[list[int]],
+) -> dict[str, Any]:
+    return {
+        "artifact_name": artifact_name,
+        "roi0_model_path": str(Path(args.roi0_model_path).resolve()),
+        "roi1_model_path": str(Path(args.roi1_model_path).resolve()),
+        "input_names": input_names,
+        "input_sizes_wh": input_sizes_wh,
+        "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),
+        "hybrid_outputs": bool(args.hybrid_outputs),
+        "denorm_branch_outputs": bool(args.denorm_branch_outputs),
+        "postprocessed_outputs": bool(args.postprocessed_outputs),
+        "edge_head_mode": args.edge_head_mode,
+        "keep_edge_head": args.edge_head_mode == "keep",
+        "fake_3d_branch_mode": args.fake_3d_branch_mode,
+        "keep_fake_3d_branch": args.fake_3d_branch_mode == "keep",
+        "export_mode": _export_mode(args),
+    }
+
+
+def _export_torchscript(
+    model: nn.Module,
+    example_inputs: torch.Tensor | tuple[torch.Tensor, ...],
+    save_path: Path,
+    manifest: dict[str, Any],
+) -> None:
+    LOGGER.info(f"Exporting TorchScript to {save_path}")
+    traced = torch.jit.trace(model, example_inputs, strict=False)
+    extra_files = {"config.txt": json.dumps(manifest, sort_keys=True)}
+    traced.save(str(save_path), _extra_files=extra_files)
+    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_inputs: torch.Tensor | tuple[torch.Tensor, ...],
+    save_path: Path,
+    input_names: list[str],
+    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_inputs,
+                str(save_path),
+                verbose=False,
+                opset_version=opset,
+                do_constant_folding=True,
+                input_names=input_names,
+                output_names=output_names,
+                dynamic_axes=_dynamic_axes(input_names, 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 _load_and_prepare_wrapper(weights_path: str, roi_name: str, args: argparse.Namespace) -> ROISelectedOutputsWrapper:
+    LOGGER.info(f"Loading {roi_name} checkpoint from {weights_path}")
+    model, _ = load_checkpoint(weights_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)
+    return ROISelectedOutputsWrapper(
+        prepared_model,
+        roi_name,
+        include_anchor_stride=args.include_anchor_stride,
+        hybrid_outputs=args.hybrid_outputs,
+        denorm_branch_outputs=args.denorm_branch_outputs,
+        postprocessed_outputs=args.postprocessed_outputs,
+        keep_edge_head=args.edge_head_mode == "keep",
+        keep_fake_3d_branch=args.fake_3d_branch_mode == "keep",
+    )
+
+
+def _export_single_roi_artifacts(
+    wrapper: ROISelectedOutputsWrapper,
+    example_input: torch.Tensor,
+    save_dir: Path,
+    base_name: str,
+    args: argparse.Namespace,
+    input_size_wh: tuple[int, int],
+) -> None:
+    outputs = wrapper(example_input)
+    output_names = _output_names(
+        base_name,
+        include_anchor_stride=args.include_anchor_stride,
+        hybrid_outputs=args.hybrid_outputs,
+        denorm_branch_outputs=args.denorm_branch_outputs,
+        postprocessed_outputs=args.postprocessed_outputs,
+        keep_edge_head=args.edge_head_mode == "keep",
+        keep_fake_3d_branch=args.fake_3d_branch_mode == "keep",
+    )
+    manifest = _base_manifest(
+        artifact_name=base_name,
+        args=args,
+        input_names=[f"{base_name}_input"],
+        input_sizes_wh={f"{base_name}_input": list(input_size_wh)},
+        output_names=output_names,
+        output_shapes=_describe_outputs(outputs),
+    )
+    _write_manifest(save_dir / f"{base_name}.export.json", manifest)
+
+    if not args.skip_torchscript:
+        _export_torchscript(wrapper, example_input, save_dir / f"{base_name}.torchscript", manifest)
+    if not args.skip_onnx:
+        _export_onnx(
+            wrapper,
+            example_input,
+            save_dir / f"{base_name}.onnx",
+            input_names=[f"{base_name}_input"],
+            output_names=output_names,
+            manifest=manifest,
+            opset=args.opset,
+            dynamic=args.dynamic,
+            simplify=args.simplify,
+        )
+
+
+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)
+
+    common_imgsz_wh = (int(args.imgsz[0]), int(args.imgsz[1]))
+    roi0_imgsz_wh = _resolve_imgsz_wh(common_imgsz_wh, args.roi0_imgsz)
+    roi1_imgsz_wh = _resolve_imgsz_wh(common_imgsz_wh, args.roi1_imgsz)
+
+    roi0_wrapper = _load_and_prepare_wrapper(args.roi0_model_path, "roi0", args)
+    roi1_wrapper = _load_and_prepare_wrapper(args.roi1_model_path, "roi1", args)
+
+    roi0_example_input = _make_example_input(roi0_imgsz_wh)
+    roi1_example_input = _make_example_input(roi1_imgsz_wh)
+
+    if args.export_separate:
+        LOGGER.info("Exporting ROI0 and ROI1 as separate artifacts")
+        with torch.inference_mode():
+            _export_single_roi_artifacts(roi0_wrapper, roi0_example_input, save_dir, "roi0_model", args, roi0_imgsz_wh)
+            _export_single_roi_artifacts(roi1_wrapper, roi1_example_input, save_dir, "roi1_model", args, roi1_imgsz_wh)
+        return
+
+    merged_model = TwoROIMergedExportModel(
+        roi0_wrapper.model,
+        roi1_wrapper.model,
+        include_anchor_stride=args.include_anchor_stride,
+        hybrid_outputs=args.hybrid_outputs,
+        denorm_branch_outputs=args.denorm_branch_outputs,
+        postprocessed_outputs=args.postprocessed_outputs,
+        keep_edge_head=args.edge_head_mode == "keep",
+        keep_fake_3d_branch=args.fake_3d_branch_mode == "keep",
+    ).eval()
+    example_inputs = (roi0_example_input, roi1_example_input)
+
+    with torch.inference_mode():
+        outputs = merged_model(*example_inputs)
+    output_names = _output_names(
+        "roi0",
+        include_anchor_stride=args.include_anchor_stride,
+        hybrid_outputs=args.hybrid_outputs,
+        denorm_branch_outputs=args.denorm_branch_outputs,
+        postprocessed_outputs=args.postprocessed_outputs,
+        keep_edge_head=args.edge_head_mode == "keep",
+        keep_fake_3d_branch=args.fake_3d_branch_mode == "keep",
+    ) + _output_names(
+        "roi1",
+        include_anchor_stride=args.include_anchor_stride,
+        hybrid_outputs=args.hybrid_outputs,
+        denorm_branch_outputs=args.denorm_branch_outputs,
+        postprocessed_outputs=args.postprocessed_outputs,
+        keep_edge_head=args.edge_head_mode == "keep",
+        keep_fake_3d_branch=args.fake_3d_branch_mode == "keep",
+    )
+    LOGGER.info(f"Merged model output shapes: {_describe_outputs(outputs)}")
+
+    manifest = _base_manifest(
+        artifact_name="merged_model",
+        args=args,
+        input_names=["roi0_input", "roi1_input"],
+        input_sizes_wh={
+            "roi0_input": list(roi0_imgsz_wh),
+            "roi1_input": list(roi1_imgsz_wh),
+        },
+        output_names=output_names,
+        output_shapes=_describe_outputs(outputs),
+    )
+    _write_manifest(save_dir / "merged_model.export.json", manifest)
+
+    if not args.skip_torchscript:
+        _export_torchscript(merged_model, example_inputs, save_dir / "merged_model.torchscript", manifest)
+    if not args.skip_onnx:
+        _export_onnx(
+            merged_model,
+            example_inputs,
+            save_dir / "merged_model.onnx",
+            input_names=["roi0_input", "roi1_input"],
+            output_names=output_names,
+            manifest=manifest,
+            opset=args.opset,
+            dynamic=args.dynamic,
+            simplify=args.simplify,
+        )
+
+
+if __name__ == "__main__":
+    main()

tools/model_merging/merge_models_of_2roi_yolo26.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
+export PYTHONPATH="${PROJECT_ROOT}:${PYTHONPATH}"
+PYTHON_BIN="/deeplearning_team/ydong/dongying/miniconda/envs/pytorch1.9/bin/python"
+
+if [ ! -x "${PYTHON_BIN}" ]; then
+    echo "ERROR: pytorch1.9 python not found: ${PYTHON_BIN}"
+    exit 1
+fi
+
+"${PYTHON_BIN}" tools/model_merging/merge_models_of_2roi_yolo26.py \
+    --roi0-model-path runs/detect/mono3d_roi0_20260506_epoch99.pt \
+    --roi1-model-path runs/detect/mono3d_roi1_20260506_epoch99.pt \
+    --save-dir runs/export/train_mono3d_two_roi_20260506-drop_fake_3d_branch \
+    --imgsz 768 352 \
+    --opset 11 \
+    --edge-head-mode drop \
+    --fake-3d-branch-mode drop
+    # --postprocessed-outputs