Добавлены пропсы конвейера и стереодвижки, задействованные в прогоне

assets/conveyors (274 МБ) - ленты и угловая секция NVIDIA, на которые ссылается сцена
относительным путём. Раньше исключались как перекачиваемые, но без них сцена не
композится из коробки.

cv/ - код стереодвижков, которые вызывает control_test, без весов:
* defom-stereo - рабочий бейзлайн (DEFOM vitl, вход 480, iters 24)
* crestereo - второй движок, точнее по габаритам (MAE 23.5 против 32.8 мм)
* fast-foundationstereo - проверялся, в бейзлайн не вошёл
* circular_section.py - показатель кругового сечения, перенесён в measure_plane.py:
  выравнивает облако по СОБСТВЕННЫМ главным осям и режет на пяти высотах вдоль каждой.
  Три самодельные версии (мировые оси, одно сечение) давали хуже; результаты проверки
  на эталонной геометрии - в circular_section_results.json

Веса по-прежнему не в репозитории - источники в MODELS.md. Наборы кадров прежних
прогонов (cv/flow_*, 1.26 ГБ) исключены: это выход, а не исходники.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

cv/fast-foundationstereo/scripts/build_plugin_trt.py

@@ -0,0 +1,150 @@
+#!/usr/bin/env python3
+"""Build a TensorRT engine from the FFSGWCVolume plugin ONNX.
+
+This mirrors cpp/app/build_single_engine.cpp in Python. The ONNX parser still
+needs the custom FFSGWCVolume plugin creator registered before parsing, so the
+shared plugin library is auto-detected from cpp/build or can be passed with
+--plugin_lib.
+"""
+
+import argparse
+import ctypes
+import os
+from pathlib import Path
+
+
+PLUGIN_NAME = "FFSGWCVolume"
+PLUGIN_VERSION = "1"
+_LOADED_PLUGIN_LIBS = []
+
+
+def find_default_plugin_library() -> Path | None:
+    repo_dir = Path(__file__).resolve().parents[1]
+    candidates = [
+        repo_dir / "cpp" / "build" / "libffs_gwc_plugin.so",
+        repo_dir / "cpp" / "build" / "lib" / "libffs_gwc_plugin.so",
+        repo_dir / "cpp" / "build" / "Release" / "libffs_gwc_plugin.so",
+    ]
+    for path in candidates:
+        if path.exists():
+            return path
+    return None
+
+
+def load_plugin_library(path: str) -> None:
+    """Load an optional shared library that registers FFSGWCVolume."""
+    lib = ctypes.CDLL(path, mode=ctypes.RTLD_GLOBAL)
+    _LOADED_PLUGIN_LIBS.append(lib)
+
+    # The current C++ code registers via ffs_depth::registerFFSGWCPlugin().
+    # A loadable Python plugin library should expose an extern "C" wrapper with
+    # one of these names so ctypes can call it without C++ name mangling.
+    for symbol in ("ffs_register_gwc_plugin", "registerFFSGWCPlugin"):
+        try:
+            fn = getattr(lib, symbol)
+        except AttributeError:
+            continue
+        fn.restype = ctypes.c_bool
+        if not fn():
+            raise RuntimeError(f"{symbol}() returned false for {path}")
+        return
+
+    # Some TensorRT plugin libraries register creators during library load. That
+    # is not true for this repo's current static C++ helper, but allow it here.
+
+
+def find_plugin_creator(trt) -> bool:
+    registry = trt.get_plugin_registry()
+    creator = registry.get_plugin_creator(PLUGIN_NAME, PLUGIN_VERSION, "")
+    return creator is not None
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Build a TensorRT engine from an ONNX graph containing FFSGWCVolume."
+    )
+    parser.add_argument("plugin_onnx", type=Path, help="Path to plugin ONNX file")
+    parser.add_argument("output_engine", type=Path, help="Path to write TensorRT engine")
+    parser.add_argument(
+        "--plugin_lib",
+        type=Path,
+        default=None,
+        help=(
+            "Shared library that registers FFSGWCVolume. Defaults to "
+            "cpp/build/libffs_gwc_plugin.so when present. Pure Python cannot "
+            "provide this repo's CUDA plugin implementation."
+        ),
+    )
+    parser.add_argument(
+        "--fp32",
+        action="store_true",
+        help="Disable FP16 builder flag. Default allows FP16 when supported.",
+    )
+    parser.add_argument(
+        "--workspace-mb",
+        type=int,
+        default=4096,
+        help="TensorRT workspace memory limit in MiB.",
+    )
+    return parser.parse_args()
+
+
+def main() -> int:
+    args = parse_args()
+    if not args.plugin_onnx.exists():
+        raise FileNotFoundError(f"ONNX file does not exist: {args.plugin_onnx}")
+    args.output_engine.parent.mkdir(parents=True, exist_ok=True)
+
+    import tensorrt as trt
+
+    logger = trt.Logger(trt.Logger.INFO)
+    trt.init_libnvinfer_plugins(logger, "")
+
+    plugin_lib = args.plugin_lib or find_default_plugin_library()
+    if plugin_lib is not None:
+        if not plugin_lib.exists():
+            raise FileNotFoundError(f"Plugin library does not exist: {plugin_lib}")
+        load_plugin_library(str(plugin_lib))
+
+    if not find_plugin_creator(trt):
+        raise RuntimeError(
+            f"{PLUGIN_NAME} plugin creator is not registered. "
+            "Build/load a shared library for cpp/src/gwc_volume_plugin.cpp and pass "
+            "--plugin_lib, or use cpp/build/ffs_build_single_engine."
+        )
+
+    explicit_batch = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
+    builder = trt.Builder(logger)
+    network = builder.create_network(explicit_batch)
+    parser = trt.OnnxParser(network, logger)
+
+    parsed = False
+    if hasattr(parser, "parse_from_file"):
+        parsed = parser.parse_from_file(str(args.plugin_onnx))
+    else:
+        parsed = parser.parse(args.plugin_onnx.read_bytes())
+    if not parsed:
+        for i in range(parser.num_errors):
+            print(parser.get_error(i))
+        raise RuntimeError(f"failed to parse ONNX: {args.plugin_onnx}")
+
+    config = builder.create_builder_config()
+    config.set_memory_pool_limit(
+        trt.MemoryPoolType.WORKSPACE, int(args.workspace_mb) * 1024 * 1024
+    )
+    if not args.fp32 and builder.platform_has_fast_fp16:
+        config.set_flag(trt.BuilderFlag.FP16)
+
+    serialized = builder.build_serialized_network(network, config)
+    if serialized is None:
+        raise RuntimeError("build_serialized_network failed")
+
+    args.output_engine.write_bytes(bytes(serialized))
+    precision = "FP32" if args.fp32 else "FP16 allowed"
+    print(f"Built engine: {args.output_engine}")
+    print(f"Precision: {precision}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

cv/fast-foundationstereo/scripts/make_onnx.py

@@ -0,0 +1,88 @@
+import warnings, argparse, logging, os, sys,zipfile
+os.environ['TORCH_COMPILE_DISABLE'] = '1'
+os.environ['TORCHDYNAMO_DISABLE'] = '1'
+code_dir = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(f'{code_dir}/../')
+import omegaconf, yaml, torch,pdb
+from omegaconf import OmegaConf
+from core.foundation_stereo import FastFoundationStereo, TrtFeatureRunner, TrtPostRunner, build_gwc_volume_triton
+import Utils as U
+
+
+class FoundationStereoOnnx(FastFoundationStereo):
+    def __init__(self, args):
+        super().__init__(args)
+
+    @torch.no_grad()
+    def forward(self, left, right):
+        """ Removes extra outputs and hyper-parameters """
+        with torch.amp.autocast('cuda', enabled=True, dtype=U.AMP_DTYPE):
+            disp = FastFoundationStereo.forward(self, left, right, iters=self.args.valid_iters, test_mode=True, optimize_build_volume=False)
+        return disp
+
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    code_dir = os.path.dirname(os.path.abspath(__file__))
+    parser.add_argument('--model_dir', type=str, default=f'{code_dir}/../weights/model_best_bp2_serialize.pth')
+    parser.add_argument('--save_path', type=str, default=f'/home/bowen/debug/', help='Path to save results.')
+    parser.add_argument('--height', type=int, default=448)
+    parser.add_argument('--width', type=int, default=640)
+    parser.add_argument('--valid_iters', type=int, default=8, help='number of flow-field updates during forward pass')
+    parser.add_argument('--corr_levels', type=int, default=2, help="number of levels in the correlation pyramid")
+    parser.add_argument('--mixed_precision', default=True, action='store_true', help='use mixed precision')
+    parser.add_argument('--corr_radius', type=int, default=4, help="width of the correlation pyramid")
+    parser.add_argument('--n_downsample', type=int, default=2, help="resolution of the disparity field (1/2^K)")
+    parser.add_argument('--n_gru_layers', type=int, default=1, help="number of hidden GRU levels")
+    parser.add_argument('--max_disp', type=int, default=192, help="max disp of geometry encoding volume")
+    parser.add_argument('--low_memory', type=int, default=1, help='reduce memory usage')
+    args = parser.parse_args()
+    os.makedirs(os.path.dirname(args.save_path), exist_ok=True)
+
+    torch.autograd.set_grad_enabled(False)
+
+    model = torch.load(args.model_dir, map_location='cpu', weights_only=False)
+    model.args.max_disp = args.max_disp
+    model.args.valid_iters = args.valid_iters
+    model.cuda().eval()
+
+    feature_runner = TrtFeatureRunner(model)
+    post_runner = TrtPostRunner(model)
+
+    feature_runner.cuda().eval()
+    post_runner.cuda().eval()
+    assert args.height % 32 == 0 and args.width % 32 == 0, "height and width must be divisible by 32"
+    left_img = torch.randn(1, 3, args.height, args.width).cuda().float()*255
+    right_img = torch.randn(1, 3, args.height, args.width).cuda().float()*255
+
+    torch.onnx.export(
+        feature_runner,
+        (left_img, right_img),
+        args.save_path+'/feature_runner.onnx',
+        opset_version=17,
+        input_names = ['left', 'right'],
+        output_names = ['features_left_04', 'features_left_08', 'features_left_16', 'features_left_32', 'features_right_04', 'stem_2x'],
+        do_constant_folding=True,
+        dynamo=False,
+    )
+
+    features_left_04, features_left_08, features_left_16, features_left_32, features_right_04, stem_2x = feature_runner(left_img, right_img)
+    gwc_volume = build_gwc_volume_triton(features_left_04.half(), features_right_04.half(), args.max_disp//4, model.cv_group)
+    disp = post_runner(features_left_04.float(), features_left_08.float(), features_left_16.float(), features_left_32.float(), features_right_04.float(), stem_2x.float(), gwc_volume.float())
+
+    torch.onnx.export(
+        post_runner,
+        (features_left_04, features_left_08, features_left_16, features_left_32, features_right_04, stem_2x, gwc_volume),
+        args.save_path+'/post_runner.onnx',
+        opset_version=17,
+        input_names = ['features_left_04', 'features_left_08', 'features_left_16', 'features_left_32', 'features_right_04', 'stem_2x', 'gwc_volume'],
+        output_names = ['disp'],
+        do_constant_folding=True,
+        dynamo=False,
+    )
+
+    with open(f'{args.save_path}/onnx.yaml', 'w') as f:
+      cfg = OmegaConf.to_container(model.args)
+      cfg['image_size'] = [args.height, args.width]
+      yaml.safe_dump(cfg, f)

cv/fast-foundationstereo/scripts/make_plugin_onnx.py

@@ -0,0 +1,160 @@
+import argparse
+import logging
+import os
+import sys
+
+os.environ['TORCH_COMPILE_DISABLE'] = '1'
+os.environ['TORCHDYNAMO_DISABLE'] = '1'
+
+code_dir = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(f'{code_dir}/../')
+
+
+def build_parser():
+    parser = argparse.ArgumentParser(
+        description='Export Fast-FoundationStereo as one ONNX with an FFSGWCVolume TensorRT plugin node')
+    parser.add_argument('--model_dir', type=str,
+                        default=f'{code_dir}/../weights/23-36-37/model_best_bp2_serialize.pth')
+    parser.add_argument('--save_path', type=str, default=f'{code_dir}/../output_plugin_onnx')
+    parser.add_argument('--height', type=int, default=608)
+    parser.add_argument('--width', type=int, default=960)
+    parser.add_argument('--valid_iters', type=int, default=8)
+    parser.add_argument('--max_disp', type=int, default=192)
+    parser.add_argument('--onnx_name', type=str, default='fast_foundationstereo_plugin.onnx')
+    return parser
+
+
+if any(arg in ('-h', '--help') for arg in sys.argv[1:]):
+    build_parser().print_help()
+    sys.exit(0)
+
+import torch
+import torch.nn as nn
+import yaml
+from omegaconf import OmegaConf
+from torch.onnx import symbolic_helper
+
+from core.foundation_stereo import TrtFeatureRunner, TrtPostRunner
+
+
+class FFSGWCVolumeOp(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, features_left_04, features_right_04, max_disp, cv_group, normalize):
+        # ONNX export only needs a tensor with the correct static shape here.
+        # symbolic() emits the TensorRT plugin node that computes the real volume.
+        batch, _, height, width = features_left_04.shape
+        return features_left_04.new_zeros(
+            (batch, int(cv_group), int(max_disp), height, width)
+        )
+
+    @staticmethod
+    def symbolic(g, features_left_04, features_right_04, max_disp, cv_group, normalize):
+        def as_int(value):
+            if isinstance(value, int):
+                return value
+            return symbolic_helper._parse_arg(value, 'i')
+
+        max_disp = as_int(max_disp)
+        cv_group = as_int(cv_group)
+        normalize = as_int(normalize)
+        out = g.op(
+            'FFSGWCVolume',
+            features_left_04,
+            features_right_04,
+            max_disp_i=int(max_disp),
+            cv_group_i=int(cv_group),
+            normalize_i=int(normalize),
+        )
+        sizes = features_left_04.type().sizes()
+        if sizes is not None and len(sizes) == 4:
+            out.setType(features_left_04.type().with_sizes(
+                [sizes[0], int(cv_group), int(max_disp), sizes[2], sizes[3]]))
+        return out
+
+
+class FastFoundationStereoPluginOnnx(nn.Module):
+    def __init__(self, model, max_disp_levels, cv_group, normalize):
+        super().__init__()
+        self.feature_runner = TrtFeatureRunner(model)
+        self.post_runner = TrtPostRunner(model)
+        self.max_disp_levels = int(max_disp_levels)
+        self.cv_group = int(cv_group)
+        self.normalize = int(bool(normalize))
+
+    @torch.no_grad()
+    def forward(self, left, right):
+        features_left_04, features_left_08, features_left_16, features_left_32, features_right_04, stem_2x = (
+            self.feature_runner(left, right)
+        )
+        gwc_volume = FFSGWCVolumeOp.apply(
+            features_left_04,
+            features_right_04,
+            self.max_disp_levels,
+            self.cv_group,
+            self.normalize,
+        )
+        disp = self.post_runner(
+            features_left_04.float(),
+            features_left_08.float(),
+            features_left_16.float(),
+            features_left_32.float(),
+            features_right_04.float(),
+            stem_2x.float(),
+            gwc_volume.float(),
+        )
+        return disp
+
+
+if __name__ == '__main__':
+    args = build_parser().parse_args()
+
+    logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
+    assert args.height % 32 == 0 and args.width % 32 == 0, 'height and width must be divisible by 32'
+    os.makedirs(args.save_path, exist_ok=True)
+    torch.autograd.set_grad_enabled(False)
+
+    logging.info('Loading model: %s', args.model_dir)
+    model = torch.load(args.model_dir, map_location='cpu', weights_only=False)
+    model.args.max_disp = args.max_disp
+    model.args.valid_iters = args.valid_iters
+    model.cuda().eval()
+
+    cv_group = int(getattr(model, 'cv_group', getattr(model.args, 'cv_group', 8)))
+    normalize = bool(getattr(model.args, 'normalize', True))
+    wrapper = FastFoundationStereoPluginOnnx(
+        model,
+        max_disp_levels=args.max_disp // 4,
+        cv_group=cv_group,
+        normalize=normalize,
+    ).cuda().eval()
+
+    left = torch.randn(1, 3, args.height, args.width, device='cuda').float() * 255
+    right = torch.randn(1, 3, args.height, args.width, device='cuda').float() * 255
+
+    onnx_name = args.onnx_name if args.onnx_name.endswith('.onnx') else f'{args.onnx_name}.onnx'
+    onnx_path = os.path.join(args.save_path, onnx_name)
+    logging.info('Exporting plugin ONNX: %s', onnx_path)
+
+    torch.onnx.export(
+        wrapper,
+        (left, right),
+        onnx_path,
+        opset_version=17,
+        input_names=['left', 'right'],
+        output_names=['disp'],
+        do_constant_folding=True,
+        dynamo=False,
+    )
+
+    cfg = OmegaConf.to_container(model.args)
+    cfg['image_size'] = [args.height, args.width]
+    cfg['cv_group'] = cv_group
+    cfg['normalize'] = normalize
+    with open(os.path.join(args.save_path, 'onnx.yaml'), 'w') as f:
+        yaml.safe_dump(cfg, f)
+
+    logging.info('ONNX model: %s', onnx_path)
+    logging.info('Config    : %s', os.path.join(args.save_path, 'onnx.yaml'))
+    logging.info('Build with:')
+    logging.info('  cpp/build/ffs_build_single_engine %s %s',
+                 onnx_path, os.path.join(args.save_path, 'fast_foundationstereo.engine'))

cv/fast-foundationstereo/scripts/make_single_onnx.py

@@ -0,0 +1,222 @@
+# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Export Fast FoundationStereo as a **single** ONNX model.
+
+Unlike make_onnx.py (which splits into feature_runner + post_runner with a
+Triton GWC kernel in between), this script produces one self-contained ONNX
+that can be converted to a single TensorRT engine via trtexec.
+
+Key design choices:
+  - The GWC and concat cost volumes are built with ONNX-compatible ops
+    (pad + slice + stack).  The upstream pytorch1 variants use
+    Tensor.unfold / torch.flip which the ONNX exporter cannot handle.
+  - ImageNet normalization is STRIPPED from the model so that it can be
+    applied externally (e.g. via Isaac ROS ImageNormalizeNode).  The ONNX
+    model expects **pre-normalized** float inputs:
+        pixel = (pixel_0_255 - mean) / std
+        mean  = [123.675, 116.28, 103.53]   (ImageNet, in 0-255 scale)
+        std   = [ 58.395, 57.12, 57.375]
+  - Inputs:  left_image  (1, 3, H, W)  float32, ImageNet-normalised
+             right_image (1, 3, H, W)  float32, ImageNet-normalised
+  - Output:  disparity   (1, 1, H, W)  float32
+
+Usage:
+  python make_single_onnx.py \\
+      --model_dir ../weights/model_best_bp2_serialize.pth \\
+      --save_path ./output_single_onnx --height 480 --width 640
+
+  # Then build a TensorRT engine:
+  trtexec --onnx=./output_single_onnx/fast_foundationstereo.onnx \\
+      --saveEngine=fast_foundationstereo.engine --fp16
+"""
+
+import argparse
+import logging
+import os
+import sys
+
+os.environ['TORCH_COMPILE_DISABLE'] = '1'
+os.environ['TORCHDYNAMO_DISABLE'] = '1'
+
+code_dir = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(f'{code_dir}/../')
+
+import yaml
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from omegaconf import OmegaConf
+import core.foundation_stereo as _fs_module
+from core.foundation_stereo import FastFoundationStereo
+
+
+# ---------------------------------------------------------------------------
+# ONNX-compatible cost-volume builders
+#
+# The upstream *_optimized_pytorch1 variants use Tensor.unfold + torch.flip
+# which the ONNX tracer cannot export.  These replacements build the
+# disparity-shifted target volume with an explicit loop over disparities
+# using only F.pad, slicing, and torch.stack — all fully ONNX-exportable.
+# ---------------------------------------------------------------------------
+
+def _build_gwc_volume_onnx(refimg_fea, targetimg_fea, maxdisp,
+                           num_groups, normalize=True):
+    dtype = refimg_fea.dtype
+    B, C, H, W = refimg_fea.shape
+    channels_per_group = C // num_groups
+
+    ref_volume = refimg_fea.unsqueeze(2).expand(B, C, maxdisp, H, W)
+
+    shifted = [
+        F.pad(targetimg_fea, (d, 0, 0, 0), 'constant', 0.0)[:, :, :, :W]
+        for d in range(maxdisp)
+    ]
+    target_volume = torch.stack(shifted, dim=2)
+
+    ref_volume = ref_volume.view(B, num_groups, channels_per_group,
+                                 maxdisp, H, W)
+    target_volume = target_volume.view(B, num_groups, channels_per_group,
+                                       maxdisp, H, W)
+
+    if normalize:
+        ref_volume = F.normalize(ref_volume.float(), dim=2).to(dtype)
+        target_volume = F.normalize(target_volume.float(), dim=2).to(dtype)
+
+    return (ref_volume * target_volume).sum(dim=2).contiguous()
+
+
+def _build_concat_volume_onnx(refimg_fea, targetimg_fea, maxdisp):
+    B, C, H, W = refimg_fea.shape
+
+    ref_volume = refimg_fea.unsqueeze(2).expand(B, C, maxdisp, H, W)
+
+    shifted = [
+        F.pad(targetimg_fea, (d, 0, 0, 0), 'constant', 0.0)[:, :, :, :W]
+        for d in range(maxdisp)
+    ]
+    target_volume = torch.stack(shifted, dim=2)
+
+    return torch.cat((ref_volume, target_volume), dim=1).contiguous()
+
+
+# ---------------------------------------------------------------------------
+
+class FastFoundationStereoSingleOnnx(nn.Module):
+    """Thin wrapper that calls the full model with ONNX-compatible settings.
+
+    Before ONNX tracing the caller must monkey-patch:
+      - normalize_image        → identity  (normalization done externally)
+      - build_gwc_volume_*     → _build_gwc_volume_onnx
+      - build_concat_volume_*  → _build_concat_volume_onnx
+    """
+
+    def __init__(self, model: FastFoundationStereo):
+        super().__init__()
+        self.model = model
+
+    @torch.no_grad()
+    def forward(self, left_image, right_image):
+        return self.model.forward(
+            left_image, right_image,
+            iters=self.model.args.valid_iters,
+            test_mode=True,
+            optimize_build_volume='pytorch1',
+        )
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(
+        description='Export Fast FoundationStereo as a single ONNX model')
+    parser.add_argument(
+        '--model_dir', type=str,
+        default=f'{code_dir}/../weights/model_best_bp2_serialize.pth',
+        help='Path to the serialized .pth model')
+    parser.add_argument(
+        '--save_path', type=str,
+        default=f'{code_dir}/output_single_onnx',
+        help='Directory to save the ONNX model and config')
+    parser.add_argument('--height', type=int, default=480)
+    parser.add_argument('--width', type=int, default=640)
+    parser.add_argument('--valid_iters', type=int, default=8,
+                        help='GRU refinement iterations')
+    parser.add_argument('--max_disp', type=int, default=192,
+                        help='Maximum disparity (in pixels at full resolution)')
+    parser.add_argument('--onnx_name', type=str, default='fast_foundationstereo',
+                        help='Base name for the saved ONNX file (without .onnx extension)')
+    args = parser.parse_args()
+
+    logging.basicConfig(
+        level=logging.INFO,
+        format='%(asctime)s %(levelname)s %(message)s')
+
+    assert args.height % 32 == 0 and args.width % 32 == 0, \
+        'height and width must be divisible by 32'
+
+    os.makedirs(args.save_path, exist_ok=True)
+    torch.autograd.set_grad_enabled(False)
+
+    if not os.path.isfile(args.model_dir):
+        raise FileNotFoundError(f'Model file not found: {args.model_dir}')
+
+    logging.info(f'Loading model from {args.model_dir}')
+    model = torch.load(args.model_dir, map_location='cpu', weights_only=False)
+    model.args.max_disp = args.max_disp
+    model.args.valid_iters = args.valid_iters
+    model.args.mixed_precision = False
+    model.cuda().eval()
+
+    wrapper = FastFoundationStereoSingleOnnx(model)
+    wrapper.cuda().eval()
+
+    left_img = torch.randn(1, 3, args.height, args.width, device='cuda')
+    right_img = torch.randn(1, 3, args.height, args.width, device='cuda')
+
+    onnx_name = args.onnx_name if args.onnx_name.endswith('.onnx') else f'{args.onnx_name}.onnx'
+    onnx_path = os.path.join(args.save_path, onnx_name)
+    logging.info(f'Exporting ONNX ({args.height}x{args.width}) → {onnx_path}')
+
+    # Monkey-patch non-ONNX-exportable functions before tracing
+    _fs_module.normalize_image = lambda img: img
+    _fs_module.build_gwc_volume_optimized_pytorch1 = _build_gwc_volume_onnx
+    _fs_module.build_concat_volume_optimized_pytorch1 = _build_concat_volume_onnx
+
+    torch.onnx.export(
+        wrapper,
+        (left_img, right_img),
+        onnx_path,
+        opset_version=17,
+        input_names=['left_image', 'right_image'],
+        output_names=['disparity'],
+        do_constant_folding=True,
+    )
+
+    cfg = OmegaConf.to_container(model.args)
+    cfg['image_size'] = [args.height, args.width]
+    config_name = os.path.splitext(onnx_name)[0] + '.yaml'
+    config_path = os.path.join(args.save_path, config_name)
+    with open(config_path, 'w') as f:
+        yaml.safe_dump(cfg, f)
+
+    logging.info(f'ONNX model  : {onnx_path}')
+    logging.info(f'Config      : {config_path}')
+    logging.info(f'Resolution  : {args.height} x {args.width}')
+    logging.info(
+        f'Build TRT engine:\n'
+        f'  trtexec --onnx={onnx_path} '
+        f'--saveEngine={args.save_path}/{os.path.splitext(onnx_name)[0]}.engine --fp16')

cv/fast-foundationstereo/scripts/profile_memory.py

@@ -0,0 +1,75 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import os,sys
+code_dir = os.path.dirname(os.path.realpath(__file__))
+sys.path.append(f'{code_dir}/../')
+from omegaconf import OmegaConf
+from core.utils.utils import InputPadder
+import argparse, torch, logging, yaml, time
+import numpy as np
+from Utils import AMP_DTYPE, set_logging_format, set_seed
+
+
+if __name__=="__main__":
+  code_dir = os.path.dirname(os.path.realpath(__file__))
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--model_dir', default=f'{code_dir}/../weights/23-36-37/model_best_bp2_serialize.pth', type=str)
+  parser.add_argument('--hiera', default=0, type=int)
+  parser.add_argument('--valid_iters', type=int, default=8, help='number of flow-field updates during forward pass')
+  parser.add_argument('--max_disp', type=int, default=192, help='maximum disparity')
+  parser.add_argument('--warmup', type=int, default=15, help='number of warmup iterations')
+  parser.add_argument('--total', type=int, default=30, help='total number of iterations')
+  args = parser.parse_args()
+
+  set_logging_format()
+  set_seed(0)
+  torch.backends.cudnn.benchmark = True
+  torch.autograd.set_grad_enabled(False)
+
+  with open(f'{os.path.dirname(args.model_dir)}/cfg.yaml', 'r') as ff:
+    cfg:dict = yaml.safe_load(ff)
+  for k in args.__dict__:
+    if args.__dict__[k] is not None:
+      cfg[k] = args.__dict__[k]
+  args = OmegaConf.create(cfg)
+  model = torch.load(args.model_dir, map_location='cpu', weights_only=False)
+  model.args.valid_iters = args.valid_iters
+  model.args.max_disp = args.max_disp
+  model.cuda().eval()
+
+  H, W = 480, 640
+  img0 = torch.randint(0, 256, (1, 3, H, W), dtype=torch.float32).cuda()
+  img1 = torch.randint(0, 256, (1, 3, H, W), dtype=torch.float32).cuda()
+  padder = InputPadder(img0.shape, divis_by=32, force_square=False)
+  img0, img1 = padder.pad(img0, img1)
+
+  logging.info(f"Image size: {H}x{W}, warmup: {args.warmup}, total: {args.total}")
+
+  times = []
+  peak_memories = []
+  with torch.amp.autocast('cuda', enabled=True, dtype=AMP_DTYPE):
+    for i in range(args.total):
+      torch.cuda.reset_peak_memory_stats()
+      torch.cuda.synchronize()
+      t0 = time.perf_counter()
+      disp = model.forward(img0, img1, iters=args.valid_iters, test_mode=True, optimize_build_volume='triton')
+      torch.cuda.synchronize()
+      elapsed = time.perf_counter() - t0
+      peak_mem = torch.cuda.max_memory_allocated() / (1024 ** 2)
+      times.append(elapsed)
+      peak_memories.append(peak_mem)
+      logging.info(f"Iter {i:2d}: {elapsed*1000:.1f} ms, peak mem: {peak_mem:.1f} MB {'(warmup)' if i < args.warmup else ''}")
+
+  measure_times = times[args.warmup:]
+  measure_mems = peak_memories[args.warmup:]
+  avg = np.mean(measure_times) * 1000
+  avg_mem = np.mean(measure_mems)
+  max_mem = np.max(measure_mems)
+  logging.info(f"Vanilla Pytorch speed average (after warmup): {avg:.1f}[ms] over {len(measure_times)} iters")
+  logging.info(f"Peak GPU memory (after warmup): avg {avg_mem:.1f} MB, max {max_mem:.1f} MB")

cv/fast-foundationstereo/scripts/profile_speed.py

@@ -0,0 +1,67 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import os,sys
+code_dir = os.path.dirname(os.path.realpath(__file__))
+sys.path.append(f'{code_dir}/../')
+from omegaconf import OmegaConf
+from core.utils.utils import InputPadder
+import argparse, torch, logging, yaml, time
+import numpy as np
+from Utils import AMP_DTYPE, set_logging_format, set_seed
+
+
+if __name__=="__main__":
+  code_dir = os.path.dirname(os.path.realpath(__file__))
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--model_dir', default=f'{code_dir}/../weights/23-36-37/model_best_bp2_serialize.pth', type=str)
+  parser.add_argument('--hiera', default=0, type=int)
+  parser.add_argument('--valid_iters', type=int, default=8, help='number of flow-field updates during forward pass')
+  parser.add_argument('--max_disp', type=int, default=192, help='maximum disparity')
+  parser.add_argument('--warmup', type=int, default=15, help='number of warmup iterations')
+  parser.add_argument('--total', type=int, default=30, help='total number of iterations')
+  args = parser.parse_args()
+
+  set_logging_format()
+  set_seed(0)
+  torch.backends.cudnn.benchmark = True
+  torch.autograd.set_grad_enabled(False)
+
+  with open(f'{os.path.dirname(args.model_dir)}/cfg.yaml', 'r') as ff:
+    cfg:dict = yaml.safe_load(ff)
+  for k in args.__dict__:
+    if args.__dict__[k] is not None:
+      cfg[k] = args.__dict__[k]
+  args = OmegaConf.create(cfg)
+  model = torch.load(args.model_dir, map_location='cpu', weights_only=False)
+  model.args.valid_iters = args.valid_iters
+  model.args.max_disp = args.max_disp
+  model.cuda().eval()
+
+  H, W = 480, 640
+  img0 = torch.randint(0, 256, (1, 3, H, W), dtype=torch.float32).cuda()
+  img1 = torch.randint(0, 256, (1, 3, H, W), dtype=torch.float32).cuda()
+  padder = InputPadder(img0.shape, divis_by=32, force_square=False)
+  img0, img1 = padder.pad(img0, img1)
+
+  logging.info(f"Image size: {H}x{W}, warmup: {args.warmup}, total: {args.total}")
+
+  times = []
+  with torch.amp.autocast('cuda', enabled=True, dtype=AMP_DTYPE):
+    for i in range(args.total):
+      torch.cuda.synchronize()
+      t0 = time.perf_counter()
+      disp = model.forward(img0, img1, iters=args.valid_iters, test_mode=True, optimize_build_volume='triton')
+      torch.cuda.synchronize()
+      elapsed = time.perf_counter() - t0
+      times.append(elapsed)
+      logging.info(f"Iter {i:2d}: {elapsed*1000:.1f} ms {'(warmup)' if i < args.warmup else ''}")
+
+  measure_times = times[args.warmup:]
+  avg = np.mean(measure_times) * 1000
+  logging.info(f"Vanilla Pytorch speed average (after warmup): {avg:.1f}[ms] over {len(measure_times)} iters")

cv/fast-foundationstereo/scripts/profile_speed_tensorrt.py

@@ -0,0 +1,66 @@
+import os,sys
+code_dir = os.path.dirname(os.path.realpath(__file__))
+sys.path.append(f'{code_dir}/../')
+from omegaconf import OmegaConf
+import argparse, torch, logging, yaml, time
+import numpy as np
+from Utils import set_logging_format, set_seed
+from core.foundation_stereo import TrtRunner
+
+
+def resolve_onnx_cfg_path(onnx_dir: str) -> str:
+  direct = os.path.join(onnx_dir, 'onnx.yaml')
+  if os.path.exists(direct):
+    return direct
+  parent = os.path.join(os.path.dirname(onnx_dir), 'onnx.yaml')
+  if os.path.exists(parent):
+    return parent
+  raise FileNotFoundError(f"onnx.yaml not found in {onnx_dir} or its parent directory")
+
+
+if __name__=="__main__":
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--model_dir', default=None, type=str)
+  parser.add_argument('--hiera', default=None, type=int)
+  parser.add_argument('--valid_iters', type=int, default=None, help='number of flow-field updates during forward pass (default: from onnx.yaml)')
+  parser.add_argument('--max_disp', type=int, default=None, help='maximum disparity (default: from onnx.yaml)')
+  parser.add_argument('--warmup', type=int, default=15, help='number of warmup iterations')
+  parser.add_argument('--total', type=int, default=30, help='total number of iterations')
+  parser.add_argument('--build_volume_backend', default=None, choices=['pytorch1', 'triton'], help='backend for cost-volume build (default: from onnx.yaml)')
+  parser.add_argument('--onnx_dir', default=f'{code_dir}/../output', type=str, help='directory containing TensorRT engines and onnx.yaml')
+  args = parser.parse_args()
+
+  set_logging_format()
+  set_seed(0)
+  torch.backends.cudnn.benchmark = True
+  torch.autograd.set_grad_enabled(False)
+
+  cfg_path = resolve_onnx_cfg_path(args.onnx_dir)
+  with open(cfg_path, 'r') as ff:
+    cfg:dict = yaml.safe_load(ff)
+  for k in args.__dict__:
+    if args.__dict__[k] is not None:
+      cfg[k] = args.__dict__[k]
+  args = OmegaConf.create(cfg)
+
+  model = TrtRunner(args, args.onnx_dir+'/feature_runner.engine', args.onnx_dir+'/post_runner.engine')
+
+  H, W = int(args.image_size[0]), int(args.image_size[1])
+  img0 = torch.randint(0, 256, (1, 3, H, W), dtype=torch.float32).cuda()
+  img1 = torch.randint(0, 256, (1, 3, H, W), dtype=torch.float32).cuda()
+
+  logging.info(f"TensorRT image size: {H}x{W}, warmup: {args.warmup}, total: {args.total}")
+
+  times = []
+  for i in range(args.total):
+    torch.cuda.synchronize()
+    t0 = time.perf_counter()
+    _ = model.forward(img0, img1)
+    torch.cuda.synchronize()
+    elapsed = time.perf_counter() - t0
+    times.append(elapsed)
+    logging.info(f"Iter {i:2d}: {elapsed*1000:.1f} ms {'(warmup)' if i < args.warmup else ''}")
+
+  measure_times = times[args.warmup:]
+  avg = np.mean(measure_times) * 1000
+  logging.info(f"TensorRT speed average (after warmup): {avg:.1f}[ms] over {len(measure_times)} iters")

cv/fast-foundationstereo/scripts/run_demo.py

@@ -0,0 +1,150 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import os,sys
+code_dir = os.path.dirname(os.path.realpath(__file__))
+sys.path.append(f'{code_dir}/../')
+from omegaconf import OmegaConf
+from core.utils.utils import InputPadder
+import argparse, torch, imageio, logging, yaml
+import numpy as np
+from Utils import (
+    AMP_DTYPE, set_logging_format, set_seed, vis_disparity,
+    depth2xyzmap, toOpen3dCloud, o3d,
+)
+import cv2
+
+
+if __name__=="__main__":
+  code_dir = os.path.dirname(os.path.realpath(__file__))
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--model_dir', default=f'{code_dir}/../weights/23-36-37/model_best_bp2_serialize.pth', type=str)
+  parser.add_argument('--left_file', default=f'{code_dir}/../demo_data/left.png', type=str)
+  parser.add_argument('--right_file', default=f'{code_dir}/../demo_data/right.png', type=str)
+  parser.add_argument('--intrinsic_file', default=f'{code_dir}/../demo_data/K.txt', type=str, help='camera intrinsic matrix and baseline file')
+  parser.add_argument('--out_dir', default='/home/bowen/debug/stereo_output', type=str)
+  parser.add_argument('--remove_invisible', default=1, type=int)
+  parser.add_argument('--denoise_cloud', default=0, type=int)
+  parser.add_argument('--denoise_nb_points', type=int, default=30, help='number of points to consider for radius outlier removal')
+  parser.add_argument('--denoise_radius', type=float, default=0.03, help='radius to use for outlier removal')
+  parser.add_argument('--scale', default=1, type=float)
+  parser.add_argument('--hiera', default=0, type=int)
+  parser.add_argument('--get_pc', type=int, default=1, help='save point cloud output')
+  parser.add_argument('--valid_iters', type=int, default=8, help='number of flow-field updates during forward pass')
+  parser.add_argument('--max_disp', type=int, default=192, help='maximum disparity')
+  parser.add_argument('--zfar', type=float, default=100, help="max depth to include in point cloud")
+  args = parser.parse_args()
+
+  set_logging_format()
+  set_seed(0)
+  torch.autograd.set_grad_enabled(False)
+
+  os.system(f'rm -rf {args.out_dir} && mkdir -p {args.out_dir}')
+
+  with open(f'{os.path.dirname(args.model_dir)}/cfg.yaml', 'r') as ff:
+    cfg:dict = yaml.safe_load(ff)
+  for k in args.__dict__:
+    if args.__dict__[k] is not None:
+      cfg[k] = args.__dict__[k]
+  args = OmegaConf.create(cfg)
+  logging.info(f"args:\n{args}")
+  model = torch.load(args.model_dir, map_location='cpu', weights_only=False)
+  model.args.valid_iters = args.valid_iters
+  model.args.max_disp = args.max_disp
+
+  model.cuda().eval()
+
+  scale = args.scale
+
+  img0 = imageio.imread(args.left_file)
+  img1 = imageio.imread(args.right_file)
+  if len(img0.shape)==2:
+    img0 = np.tile(img0[...,None], (1,1,3))
+    img1 = np.tile(img1[...,None], (1,1,3))
+  img0 = img0[...,:3]
+  img1 = img1[...,:3]
+  H,W = img0.shape[:2]
+
+  img0 = cv2.resize(img0, fx=scale, fy=scale, dsize=None)
+  img1 = cv2.resize(img1, dsize=(img0.shape[1], img0.shape[0]))
+  H,W = img0.shape[:2]
+  img0_ori = img0.copy()
+  img1_ori = img1.copy()
+  logging.info(f"img0: {img0.shape}")
+  imageio.imwrite(f'{args.out_dir}/left.png', img0)
+  imageio.imwrite(f'{args.out_dir}/right.png', img1)
+
+  img0 = torch.as_tensor(img0).cuda().float()[None].permute(0,3,1,2)
+  img1 = torch.as_tensor(img1).cuda().float()[None].permute(0,3,1,2)
+  padder = InputPadder(img0.shape, divis_by=32, force_square=False)
+  img0, img1 = padder.pad(img0, img1)
+
+  logging.info(f"Start forward, 1st time run can be slow due to compilation")
+  with torch.amp.autocast('cuda', enabled=True, dtype=AMP_DTYPE):
+    if not args.hiera:
+      disp = model.forward(img0, img1, iters=args.valid_iters, test_mode=True, optimize_build_volume='pytorch1')
+    else:
+      disp = model.run_hierachical(img0, img1, iters=args.valid_iters, test_mode=True, small_ratio=0.5)
+  logging.info("forward done")
+  disp = padder.unpad(disp.float())
+  disp = disp.data.cpu().numpy().reshape(H,W).clip(0, None)
+
+  cmap = None
+  min_val = None
+  max_val = None
+  vis = vis_disparity(disp, min_val=min_val, max_val=max_val, cmap=cmap, color_map=cv2.COLORMAP_TURBO)
+  vis = np.concatenate([img0_ori, img1_ori, vis], axis=1)
+  imageio.imwrite(f'{args.out_dir}/disp_vis.png', vis)
+  s = 1280/vis.shape[1]
+  resized_vis = cv2.resize(vis, (int(vis.shape[1]*s), int(vis.shape[0]*s)))
+  cv2.imshow('disp', resized_vis[:,:,::-1])
+  cv2.waitKey(0)
+
+  if args.remove_invisible:
+    yy,xx = np.meshgrid(np.arange(disp.shape[0]), np.arange(disp.shape[1]), indexing='ij')
+    us_right = xx-disp
+    invalid = us_right<0
+    disp[invalid] = np.inf
+
+  if args.get_pc:
+    with open(args.intrinsic_file, 'r') as f:
+      lines = f.readlines()
+      K = np.array(list(map(float, lines[0].rstrip().split()))).astype(np.float32).reshape(3,3)
+      baseline = float(lines[1])
+    K[:2] *= scale
+    depth = K[0,0]*baseline/disp
+    np.save(f'{args.out_dir}/depth_meter.npy', depth)
+    xyz_map = depth2xyzmap(depth, K)
+    pcd = toOpen3dCloud(xyz_map.reshape(-1,3), img0_ori.reshape(-1,3))
+    keep_mask = (np.asarray(pcd.points)[:,2]>0) & (np.asarray(pcd.points)[:,2]<=args.zfar)
+    keep_ids = np.arange(len(np.asarray(pcd.points)))[keep_mask]
+    pcd = pcd.select_by_index(keep_ids)
+    o3d.io.write_point_cloud(f'{args.out_dir}/cloud.ply', pcd)
+    logging.info(f"PCL saved to {args.out_dir}")
+
+    if args.denoise_cloud:
+      logging.info("[Optional step] denoise point cloud...")
+      pcd = pcd.voxel_down_sample(voxel_size=0.001)
+      cl, ind = pcd.remove_radius_outlier(nb_points=args.denoise_nb_points, radius=args.denoise_radius)
+      inlier_cloud = pcd.select_by_index(ind)
+      o3d.io.write_point_cloud(f'{args.out_dir}/cloud_denoise.ply', inlier_cloud)
+      pcd = inlier_cloud
+
+    logging.info("Visualizing point cloud. Press ESC to exit.")
+    vis = o3d.visualization.Visualizer()
+    vis.create_window()
+    vis.add_geometry(pcd)
+    vis.get_render_option().point_size = 1.0
+    vis.get_render_option().background_color = np.array([0.5, 0.5, 0.5])
+    ctr = vis.get_view_control()
+    ctr.set_front([0, 0, -1])
+    id = np.asarray(pcd.points)[:,2].argmin()
+    ctr.set_lookat(np.asarray(pcd.points)[id])
+    ctr.set_up([0, -1, 0])
+    vis.run()
+    vis.destroy_window()

cv/fast-foundationstereo/scripts/run_demo_plugin_trt.py

@@ -0,0 +1,380 @@
+#!/usr/bin/env python3
+"""Run the FFSGWCVolume-plugin TensorRT engine from Python.
+
+This is the Python equivalent of cpp/src/ffs_single_depth_inference.cpp plus
+cpp/app/main.cpp for the single-engine plugin path:
+
+  - engine inputs: left, right
+  - engine output: disp
+  - input tensors are RGB, CHW, float32, raw 0-255 values
+  - resize is aspect-ratio-preserving with right/bottom replicate padding
+  - disparity is cropped, nearest-neighbor upsampled, and scaled back to the
+    original input-image pixel units
+"""
+
+import argparse
+import logging
+import os
+import sys
+from pathlib import Path
+
+import cv2
+import imageio.v2 as imageio
+import numpy as np
+import torch
+import yaml
+
+code_dir = Path(__file__).resolve().parent
+repo_dir = code_dir.parent
+sys.path.append(str(repo_dir))
+
+from Utils import set_logging_format, set_seed, vis_disparity
+from build_plugin_trt import (
+    PLUGIN_NAME,
+    find_default_plugin_library,
+    find_plugin_creator,
+    load_plugin_library,
+)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Run Fast-FoundationStereo single TensorRT plugin engine."
+    )
+    parser.add_argument(
+        "--model_dir",
+        type=Path,
+        default=repo_dir / "engine1_plug",
+        help="Directory containing fast_foundationstereo.engine and one YAML config.",
+    )
+    parser.add_argument(
+        "--engine_file",
+        "--model_file",
+        dest="engine_file",
+        type=Path,
+        default=None,
+        help="Explicit TensorRT engine path. Defaults to <model_dir>/fast_foundationstereo.engine.",
+    )
+    parser.add_argument(
+        "--config_file",
+        type=Path,
+        default=None,
+        help="Explicit YAML config path. Defaults to the single YAML in the engine directory.",
+    )
+    parser.add_argument(
+        "--plugin_lib",
+        type=Path,
+        default=None,
+        help="Path to libffs_gwc_plugin.so. Defaults to cpp/build/libffs_gwc_plugin.so.",
+    )
+    parser.add_argument("--left_file", type=Path, default=repo_dir / "demo_data" / "left.png")
+    parser.add_argument("--right_file", type=Path, default=repo_dir / "demo_data" / "right.png")
+    parser.add_argument(
+        "--intrinsic_file",
+        type=Path,
+        default=repo_dir / "demo_data" / "K.txt",
+        help="Text file with 3x3 K on line 1 and baseline in meters on line 2.",
+    )
+    parser.add_argument("--out_dir", type=Path, default=repo_dir / "output_plugin_trt")
+    return parser.parse_args()
+
+
+def resolve_engine_path(model_dir: Path, engine_file: Path | None) -> Path:
+    path = engine_file if engine_file is not None else model_dir / "fast_foundationstereo.engine"
+    if not path.exists():
+        raise FileNotFoundError(f"TensorRT engine does not exist: {path}")
+    return path
+
+
+def resolve_config_path(engine_path: Path, model_dir: Path, config_file: Path | None) -> Path:
+    if config_file is not None:
+        if not config_file.exists():
+            raise FileNotFoundError(f"Config file does not exist: {config_file}")
+        return config_file
+
+    search_dir = model_dir if model_dir.exists() else engine_path.parent
+    yaml_files = sorted(
+        p for p in search_dir.iterdir()
+        if p.is_file() and p.suffix.lower() in (".yaml", ".yml")
+    )
+    if not yaml_files:
+        raise FileNotFoundError(f"No YAML config found in: {search_dir}")
+    if len(yaml_files) > 1:
+        names = " ".join(str(p) for p in yaml_files)
+        raise RuntimeError(
+            f"Expected exactly one YAML config in {search_dir}, found {len(yaml_files)}: {names}"
+        )
+    return yaml_files[0]
+
+
+def load_intrinsics(path: Path) -> tuple[np.ndarray, float]:
+    with path.open("r") as f:
+        lines = f.readlines()
+    if len(lines) < 2:
+        raise RuntimeError("intrinsic file must contain K on line 1 and baseline on line 2")
+    k = np.array(list(map(float, lines[0].strip().split())), dtype=np.float32).reshape(3, 3)
+    baseline = float(lines[1])
+    if not (k[0, 0] > 0.0 and baseline > 0.0):
+        raise RuntimeError("invalid focal length or baseline in intrinsic file")
+    return k, baseline
+
+
+def load_rgb(path: Path) -> np.ndarray:
+    img = imageio.imread(path)
+    if img.ndim == 2:
+        img = np.tile(img[..., None], (1, 1, 3))
+    img = img[..., :3]
+    if img.dtype != np.uint8:
+        img = np.clip(img, 0, 255).astype(np.uint8)
+    return np.ascontiguousarray(img)
+
+
+def resize_uniform_and_pad_rgb(img: np.ndarray, target_h: int, target_w: int):
+    src_h, src_w = img.shape[:2]
+    if src_h == target_h and src_w == target_w:
+        return img.copy(), target_h, target_w
+
+    scale = min(float(target_w) / src_w, float(target_h) / src_h)
+    scaled_w = max(1, int(round(src_w * scale)))
+    scaled_h = max(1, int(round(src_h * scale)))
+
+    scaled = cv2.resize(img, (scaled_w, scaled_h), interpolation=cv2.INTER_LINEAR)
+    padded = cv2.copyMakeBorder(
+        scaled,
+        0,
+        target_h - scaled_h,
+        0,
+        target_w - scaled_w,
+        cv2.BORDER_REPLICATE,
+    )
+    return np.ascontiguousarray(padded), scaled_h, scaled_w
+
+
+def tensor_from_rgb_255(img: np.ndarray) -> torch.Tensor:
+    arr = img.astype(np.float32, copy=False)
+    return torch.as_tensor(arr, device="cuda").permute(2, 0, 1).unsqueeze(0).contiguous()
+
+
+def cpp_nearest_resize(src: np.ndarray, dst_h: int, dst_w: int) -> np.ndarray:
+    src_h, src_w = src.shape
+    xs = np.rint((np.arange(dst_w, dtype=np.float32) + 0.5) * src_w / dst_w - 0.5)
+    ys = np.rint((np.arange(dst_h, dtype=np.float32) + 0.5) * src_h / dst_h - 0.5)
+    xs = np.clip(xs.astype(np.int64), 0, src_w - 1)
+    ys = np.clip(ys.astype(np.int64), 0, src_h - 1)
+    return src[ys[:, None], xs[None, :]]
+
+
+def postprocess_disparity(
+    raw_disp: np.ndarray,
+    input_h: int,
+    input_w: int,
+    model_h: int,
+    model_w: int,
+    scaled_h: int,
+    scaled_w: int,
+) -> np.ndarray:
+    disp = raw_disp.reshape(model_h, model_w).astype(np.float32, copy=False)
+    disp = np.maximum(disp, 0.0)
+
+    if input_h == model_h and input_w == model_w:
+        return np.ascontiguousarray(disp)
+
+    cropped = disp[:scaled_h, :scaled_w]
+    upsampled = cpp_nearest_resize(cropped, input_h, input_w)
+    return np.ascontiguousarray(upsampled * (float(input_w) / scaled_w), dtype=np.float32)
+
+
+def disparity_to_depth_cpp(disp: np.ndarray, fx: float, baseline_m: float) -> np.ndarray:
+    d = disp.astype(np.float32, copy=True)
+    xs = np.arange(d.shape[1], dtype=np.float32)[None, :]
+    d[(xs - d) < 0.0] = np.inf
+    with np.errstate(divide="ignore", invalid="ignore"):
+        depth = np.float32(fx * baseline_m) / d
+    return np.ascontiguousarray(depth, dtype=np.float32)
+
+
+def save_float_matrix_bin(path: Path, values: np.ndarray) -> None:
+    arr = np.ascontiguousarray(values, dtype=np.float32)
+    dims = np.array(arr.shape[:2], dtype=np.int32)
+    with path.open("wb") as f:
+        dims.tofile(f)
+        arr.tofile(f)
+
+
+def colorize_depth(depth: np.ndarray) -> np.ndarray:
+    valid = np.isfinite(depth) & (depth > 0.0)
+    if not np.any(valid):
+        return np.zeros((*depth.shape, 3), dtype=np.uint8)
+
+    safe = np.zeros_like(depth, dtype=np.float32)
+    safe[valid] = depth[valid]
+    min_val = float(safe[valid].min())
+    max_val = float(safe[valid].max())
+    if max_val <= min_val:
+        max_val = min_val + 1.0
+
+    depth_u8 = np.clip((safe - min_val) * (255.0 / (max_val - min_val)), 0, 255).astype(np.uint8)
+    colored_bgr = cv2.applyColorMap(depth_u8, cv2.COLORMAP_TURBO)
+    colored = colored_bgr[..., ::-1]
+    colored[~valid] = 0
+    return colored
+
+
+class PluginTensorRTRunner:
+    def __init__(self, engine_path: Path, plugin_lib: Path):
+        import tensorrt as trt
+
+        self.trt = trt
+        self.logger = trt.Logger(trt.Logger.WARNING)
+        trt.init_libnvinfer_plugins(self.logger, "")
+        load_plugin_library(str(plugin_lib))
+        if not find_plugin_creator(trt):
+            raise RuntimeError(f"{PLUGIN_NAME} plugin creator is not registered after loading {plugin_lib}")
+
+        self.runtime = trt.Runtime(self.logger)
+        self.engine = self.runtime.deserialize_cuda_engine(engine_path.read_bytes())
+        if self.engine is None:
+            raise RuntimeError(f"Failed to deserialize TensorRT engine: {engine_path}")
+        self.context = self.engine.create_execution_context()
+        if self.context is None:
+            raise RuntimeError(f"Failed to create execution context: {engine_path}")
+
+        names = [self.engine.get_tensor_name(i) for i in range(self.engine.num_io_tensors)]
+        missing = [name for name in ("left", "right", "disp") if name not in names]
+        if missing:
+            raise RuntimeError(
+                f"Plugin engine must expose tensors named left, right, and disp. "
+                f"Missing {missing}; found {names}"
+            )
+
+    def _torch_dtype(self, trt_dtype):
+        trt = self.trt
+        mapping = {
+            trt.DataType.FLOAT: torch.float32,
+            trt.DataType.HALF: torch.float16,
+            trt.DataType.INT32: torch.int32,
+            trt.DataType.INT8: torch.int8,
+            trt.DataType.BOOL: torch.bool,
+        }
+        if hasattr(trt.DataType, "BF16"):
+            mapping[trt.DataType.BF16] = torch.bfloat16
+        if trt_dtype not in mapping:
+            raise RuntimeError(f"Unsupported TensorRT dtype: {trt_dtype}")
+        return mapping[trt_dtype]
+
+    def infer(self, left: torch.Tensor, right: torch.Tensor) -> torch.Tensor:
+        inputs = {"left": left, "right": right}
+        for name, tensor in list(inputs.items()):
+            expected = self._torch_dtype(self.engine.get_tensor_dtype(name))
+            if tensor.dtype != expected:
+                tensor = tensor.to(expected)
+            inputs[name] = tensor.contiguous()
+            self.context.set_input_shape(name, tuple(inputs[name].shape))
+
+        out_shape = tuple(self.context.get_tensor_shape("disp"))
+        out_dtype = self._torch_dtype(self.engine.get_tensor_dtype("disp"))
+        output = torch.empty(out_shape, device="cuda", dtype=out_dtype)
+
+        self.context.set_tensor_address("left", int(inputs["left"].data_ptr()))
+        self.context.set_tensor_address("right", int(inputs["right"].data_ptr()))
+        self.context.set_tensor_address("disp", int(output.data_ptr()))
+
+        stream = torch.cuda.current_stream().cuda_stream
+        if not self.context.execute_async_v3(stream):
+            raise RuntimeError("TensorRT enqueue failed")
+        return output
+
+
+def main() -> int:
+    args = parse_args()
+    set_logging_format()
+    set_seed(0)
+    torch.autograd.set_grad_enabled(False)
+    args.out_dir.mkdir(parents=True, exist_ok=True)
+
+    engine_path = resolve_engine_path(args.model_dir, args.engine_file)
+    config_path = resolve_config_path(engine_path, args.model_dir, args.config_file)
+    plugin_lib = args.plugin_lib or find_default_plugin_library()
+    if plugin_lib is None or not plugin_lib.exists():
+        raise FileNotFoundError(
+            "Could not find libffs_gwc_plugin.so. Build cpp first or pass "
+            "--plugin_lib /path/to/libffs_gwc_plugin.so"
+        )
+
+    with config_path.open("r") as f:
+        cfg = yaml.safe_load(f)
+    model_h, model_w = [int(v) for v in cfg["image_size"]]
+
+    left = load_rgb(args.left_file)
+    right = load_rgb(args.right_file)
+    if left.shape != right.shape:
+        raise RuntimeError(f"left/right image size mismatch: {left.shape} vs {right.shape}")
+    input_h, input_w = left.shape[:2]
+
+    logging.info(f"Engine: {engine_path}")
+    logging.info(f"Plugin: {plugin_lib}")
+    logging.info(f"Config: {config_path}")
+    logging.info(f"Input images: {input_w}x{input_h}")
+    logging.info(f"Model target resolution: {model_w}x{model_h}")
+
+    left_model, scaled_h, scaled_w = resize_uniform_and_pad_rgb(left, model_h, model_w)
+    right_model, right_scaled_h, right_scaled_w = resize_uniform_and_pad_rgb(right, model_h, model_w)
+    if (scaled_h, scaled_w) != (right_scaled_h, right_scaled_w):
+        raise RuntimeError("left/right images produced different resize scales")
+    if (input_h, input_w) != (model_h, model_w):
+        logging.info(
+            f"Uniform resize+pad: {input_w}x{input_h} -> "
+            f"{scaled_w}x{scaled_h} inside {model_w}x{model_h}"
+        )
+
+    runner = PluginTensorRTRunner(engine_path, plugin_lib)
+    t_left = tensor_from_rgb_255(left_model)
+    t_right = tensor_from_rgb_255(right_model)
+
+    logging.info("Running TensorRT inference")
+    disp_raw_t = runner.infer(t_left, t_right)
+    torch.cuda.current_stream().synchronize()
+    logging.info("Inference done")
+
+    disp_raw = disp_raw_t.float().detach().cpu().numpy()
+    disp = postprocess_disparity(
+        disp_raw,
+        input_h=input_h,
+        input_w=input_w,
+        model_h=model_h,
+        model_w=model_w,
+        scaled_h=scaled_h,
+        scaled_w=scaled_w,
+    )
+
+    if not np.isfinite(disp).any():
+        raise RuntimeError("Model produced no finite disparity values")
+
+    k, baseline = load_intrinsics(args.intrinsic_file)
+    depth = disparity_to_depth_cpp(disp, float(k[0, 0]), baseline)
+
+    disparity_path = args.out_dir / "disparity.bin"
+    depth_bin_path = args.out_dir / "depth_meter.bin"
+    depth_npy_path = args.out_dir / "depth_meter.npy"
+    disp_vis_path = args.out_dir / "disp_vis.png"
+    depth_vis_path = args.out_dir / "depth_vis.png"
+
+    save_float_matrix_bin(disparity_path, disp)
+    save_float_matrix_bin(depth_bin_path, depth)
+    np.save(depth_npy_path, depth)
+
+    disp_color = vis_disparity(disp, color_map=cv2.COLORMAP_TURBO)
+    disp_vis = np.concatenate([left, right, disp_color], axis=1)
+    imageio.imwrite(disp_vis_path, disp_vis)
+    imageio.imwrite(depth_vis_path, colorize_depth(depth))
+
+    logging.info(f"Saved: {disparity_path}")
+    logging.info(f"Saved: {depth_bin_path}")
+    logging.info(f"Saved: {depth_npy_path}")
+    logging.info(f"Saved: {disp_vis_path}")
+    logging.info(f"Saved: {depth_vis_path}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

cv/fast-foundationstereo/scripts/run_demo_single_trt.py

@@ -0,0 +1,348 @@
+# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Run Fast FoundationStereo inference with the single ONNX model (or TRT engine)
+produced by make_single_onnx.py.
+
+Supports two backends:
+  - ONNX Runtime (default if --model_file points to an .onnx, or auto-detected)
+  - TensorRT     (if --model_file points to an .engine)
+
+The model expects ImageNet-normalised inputs, so this script applies
+normalisation during preprocessing.
+
+Usage:
+  # Run directly with ONNX (no trtexec step needed):
+  python run_demo_single_trt.py \
+      --model_dir ./output_single_onnx \
+      --left_file  ../demo_data/left.png \
+      --right_file ../demo_data/right.png
+
+  # Or with an explicit model file:
+  python run_demo_single_trt.py \
+      --model_dir  ./output_single_onnx \
+      --model_file ./output_single_onnx/fast_foundationstereo.onnx \
+      --left_file  ../demo_data/left.png \
+      --right_file ../demo_data/right.png
+"""
+
+import argparse
+import logging
+import os
+import sys
+
+import cv2
+import imageio
+import numpy as np
+import torch
+import yaml
+
+code_dir = os.path.dirname(os.path.realpath(__file__))
+sys.path.append(f'{code_dir}/../')
+
+from Utils import (
+    set_logging_format, set_seed, vis_disparity,
+    depth2xyzmap, toOpen3dCloud, o3d,
+)
+
+IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
+IMAGENET_STD  = np.array([0.229, 0.224, 0.225], dtype=np.float32)
+
+
+class SingleEngineTrtRunner:
+    """Minimal TensorRT runner for a single engine with named I/O."""
+
+    def __init__(self, engine_path):
+        import tensorrt as trt
+        self.trt = trt
+        self.logger = trt.Logger(trt.Logger.WARNING)
+
+        with open(engine_path, 'rb') as f:
+            self.engine = trt.Runtime(self.logger).deserialize_cuda_engine(f.read())
+        if self.engine is None:
+            raise RuntimeError(
+                f'Failed to deserialize TRT engine from {engine_path}. '
+                f'This usually means the engine was built with a different '
+                f'TensorRT version (yours: {trt.__version__}). '
+                f'Rebuild with:  trtexec --onnx=<your .onnx> '
+                f'--saveEngine={engine_path} --fp16')
+        self.context = self.engine.create_execution_context()
+
+    def _trt_to_torch_dtype(self, dt):
+        trt = self.trt
+        mapping = {
+            trt.DataType.FLOAT:  torch.float32,
+            trt.DataType.HALF:   torch.float16,
+            trt.DataType.BF16:   torch.bfloat16,
+            trt.DataType.INT32:  torch.int32,
+            trt.DataType.INT8:   torch.int8,
+            trt.DataType.BOOL:   torch.bool,
+        }
+        if dt not in mapping:
+            raise RuntimeError(f'Unsupported TRT dtype: {dt}')
+        return mapping[dt]
+
+    def __call__(self, inputs: dict) -> dict:
+        """Run inference.
+
+        Args:
+            inputs: {binding_name: torch.Tensor} for every input tensor.
+        Returns:
+            {binding_name: torch.Tensor} for every output tensor.
+        """
+        trt = self.trt
+
+        for name, tensor in inputs.items():
+            expected = self._trt_to_torch_dtype(self.engine.get_tensor_dtype(name))
+            if tensor.dtype != expected:
+                inputs[name] = tensor.to(expected)
+            if not inputs[name].is_contiguous():
+                inputs[name] = inputs[name].contiguous()
+            self.context.set_input_shape(name, tuple(inputs[name].shape))
+
+        out_names = [
+            self.engine.get_tensor_name(i)
+            for i in range(self.engine.num_io_tensors)
+            if self.engine.get_tensor_mode(self.engine.get_tensor_name(i))
+               == trt.TensorIOMode.OUTPUT
+        ]
+
+        outputs = {}
+        for name in out_names:
+            shape = tuple(self.context.get_tensor_shape(name))
+            dtype = self._trt_to_torch_dtype(self.engine.get_tensor_dtype(name))
+            outputs[name] = torch.empty(shape, device='cuda', dtype=dtype)
+
+        for name, tensor in inputs.items():
+            self.context.set_tensor_address(name, int(tensor.data_ptr()))
+        for name, tensor in outputs.items():
+            self.context.set_tensor_address(name, int(tensor.data_ptr()))
+
+        stream = torch.cuda.current_stream().cuda_stream
+        assert self.context.execute_async_v3(stream)
+
+        return outputs
+
+
+class OnnxRuntimeRunner:
+    """Run inference via ONNX Runtime (GPU if available, else CPU)."""
+
+    def __init__(self, onnx_path):
+        import onnxruntime as ort
+        providers = []
+        if 'CUDAExecutionProvider' in ort.get_available_providers():
+            providers.append('CUDAExecutionProvider')
+        providers.append('CPUExecutionProvider')
+        logging.info(f'ONNX Runtime providers: {providers}')
+        self.session = ort.InferenceSession(onnx_path, providers=providers)
+        self.input_names = [inp.name for inp in self.session.get_inputs()]
+        self.output_names = [out.name for out in self.session.get_outputs()]
+
+    def __call__(self, inputs: dict) -> dict:
+        feed = {}
+        for name in self.input_names:
+            tensor = inputs[name]
+            if isinstance(tensor, torch.Tensor):
+                tensor = tensor.cpu().float().numpy()
+            feed[name] = tensor
+        raw_outputs = self.session.run(self.output_names, feed)
+        outputs = {}
+        for name, arr in zip(self.output_names, raw_outputs):
+            outputs[name] = torch.as_tensor(arr).cuda()
+        return outputs
+
+
+def normalize_imagenet(img_uint8: np.ndarray) -> np.ndarray:
+    """Apply ImageNet normalization: (img/255 - mean) / std."""
+    return ((img_uint8.astype(np.float32) / 255.0) - IMAGENET_MEAN) / IMAGENET_STD
+
+
+def resolve_config(model_path: str) -> str:
+    """Find the YAML config matching the model file, falling back to defaults."""
+    model_dir = os.path.dirname(model_path)
+    base = os.path.splitext(os.path.basename(model_path))[0]
+    candidates = [
+        os.path.join(model_dir, f'{base}.yaml'),
+        os.path.join(model_dir, 'config.yaml'),
+        os.path.join(model_dir, 'onnx.yaml'),
+    ]
+    for p in candidates:
+        if os.path.exists(p):
+            return p
+    raise FileNotFoundError(
+        f'No .yaml config found for {model_path}. '
+        'Run make_single_onnx.py first.')
+
+
+def find_model(model_dir: str) -> str:
+    """Find an .engine or .onnx file in the directory (prefer .engine)."""
+    for ext in ('.engine', '.onnx'):
+        for f in os.listdir(model_dir):
+            if f.endswith(ext):
+                return os.path.join(model_dir, f)
+    raise FileNotFoundError(
+        f'No .engine or .onnx file found in {model_dir}. '
+        'Run make_single_onnx.py first.')
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(
+        description='Run Fast FoundationStereo with ONNX Runtime or TensorRT')
+    parser.add_argument('--model_dir', type=str,
+                        default=f'{code_dir}/output_single_onnx',
+                        help='Directory containing .onnx/.engine + config.yaml')
+    parser.add_argument('--model_file', type=str, default='',
+                        help='Explicit path to .onnx or .engine file (overrides auto-search)')
+    parser.add_argument('--left_file', type=str,
+                        default=f'{code_dir}/../demo_data/left.png')
+    parser.add_argument('--right_file', type=str,
+                        default=f'{code_dir}/../demo_data/right.png')
+    parser.add_argument('--intrinsic_file', type=str,
+                        default=f'{code_dir}/../demo_data/K.txt',
+                        help='Camera intrinsic matrix and baseline file')
+    parser.add_argument('--out_dir', type=str,
+                        default=f'{code_dir}/../output_demo')
+    parser.add_argument('--remove_invisible', type=int, default=1)
+    parser.add_argument('--denoise_cloud', type=int, default=1)
+    parser.add_argument('--denoise_nb_points', type=int, default=30)
+    parser.add_argument('--denoise_radius', type=float, default=0.03)
+    parser.add_argument('--get_pc', type=int, default=1,
+                        help='Generate and save point cloud')
+    parser.add_argument('--zfar', type=float, default=100,
+                        help='Max depth (m) to include in point cloud')
+    args = parser.parse_args()
+
+    set_logging_format()
+    set_seed(0)
+    torch.autograd.set_grad_enabled(False)
+    os.makedirs(args.out_dir, exist_ok=True)
+
+    # ── Find model and config ─────────────────────────────────────────────
+    model_path = args.model_file if args.model_file else find_model(args.model_dir)
+    cfg_path = resolve_config(model_path)
+    with open(cfg_path, 'r') as f:
+        cfg = yaml.safe_load(f)
+    target_h, target_w = cfg['image_size']
+    logging.info(f'Model target resolution: {target_h} x {target_w}')
+
+    # ── Load model (ONNX Runtime or TensorRT) ────────────────────────────
+    logging.info(f'Loading model: {model_path}')
+    if model_path.endswith('.onnx'):
+        runner = OnnxRuntimeRunner(model_path)
+    else:
+        runner = SingleEngineTrtRunner(model_path)
+
+    # ── Read images ───────────────────────────────────────────────────────
+    img0 = imageio.imread(args.left_file)
+    img1 = imageio.imread(args.right_file)
+
+    if img0.ndim == 2:
+        img0 = np.tile(img0[..., None], (1, 1, 3))
+        img1 = np.tile(img1[..., None], (1, 1, 3))
+    img0 = img0[..., :3]
+    img1 = img1[..., :3]
+
+    # ── Resize to model resolution (direct stretch) ────────────────────────
+    orig_h, orig_w = img0.shape[:2]
+    fx = target_w / orig_w
+    fy = target_h / orig_h
+
+    if fx != 1 or fy != 1:
+        logging.info(
+            f'Resizing images: {orig_h}x{orig_w} → {target_h}x{target_w} '
+            f'(fx={fx:.4f}, fy={fy:.4f})')
+        img0 = cv2.resize(img0, (target_w, target_h), interpolation=cv2.INTER_LINEAR)
+        img1 = cv2.resize(img1, (target_w, target_h), interpolation=cv2.INTER_LINEAR)
+    H, W = img0.shape[:2]
+
+    img0_ori = img0.copy()
+    img1_ori = img1.copy()
+    logging.info(f'Image size after resize: {img0.shape}')
+    imageio.imwrite(f'{args.out_dir}/left.png', img0)
+    imageio.imwrite(f'{args.out_dir}/right.png', img1)
+
+    # ── Preprocess: ImageNet normalize → NCHW float tensor ────────────────
+    img0_norm = normalize_imagenet(img0)
+    img1_norm = normalize_imagenet(img1)
+
+    t_left  = torch.as_tensor(img0_norm).cuda().float()[None].permute(0, 3, 1, 2)
+    t_right = torch.as_tensor(img1_norm).cuda().float()[None].permute(0, 3, 1, 2)
+
+    # ── Inference ─────────────────────────────────────────────────────────
+    logging.info('Running inference (first run may be slow due to TRT warmup)')
+    outputs = runner({'left_image': t_left, 'right_image': t_right})
+    disp = outputs['disparity']
+    logging.info('Inference done')
+
+    disp = disp.float().cpu().numpy().reshape(H, W).clip(0, None) * (1.0 / fx)
+
+    # ── Visualise disparity ──────────────────────────────────────────────
+    vis = vis_disparity(disp, color_map=cv2.COLORMAP_TURBO)
+    vis = np.concatenate([img0_ori, img1_ori, vis], axis=1)
+    imageio.imwrite(f'{args.out_dir}/disp_vis.png', vis)
+    s = 1280 / vis.shape[1]
+    resized_vis = cv2.resize(vis, (int(vis.shape[1] * s), int(vis.shape[0] * s)))
+    cv2.imshow('disp', resized_vis[:, :, ::-1])
+    cv2.waitKey(0)
+
+    # ── Remove invisible pixels ──────────────────────────────────────────
+    if args.remove_invisible:
+        _, xx = np.meshgrid(np.arange(H), np.arange(W), indexing='ij')
+        invalid = (xx - disp) < 0
+        disp[invalid] = np.inf
+
+    # ── Point cloud generation ───────────────────────────────────────────
+    if args.get_pc:
+        with open(args.intrinsic_file, 'r') as f:
+            lines = f.readlines()
+            K = (np.array(list(map(float, lines[0].rstrip().split())))
+                 .astype(np.float32).reshape(3, 3))
+            baseline = float(lines[1])
+        K[:2] *= np.array([fx, fy], dtype=np.float32)[:, np.newaxis]
+        depth = K[0, 0] * baseline / disp
+        np.save(f'{args.out_dir}/depth_meter.npy', depth)
+
+        xyz_map = depth2xyzmap(depth, K)
+        pcd = toOpen3dCloud(xyz_map.reshape(-1, 3), img0_ori.reshape(-1, 3))
+        pts = np.asarray(pcd.points)
+        keep = (pts[:, 2] > 0) & (pts[:, 2] <= args.zfar)
+        pcd = pcd.select_by_index(np.where(keep)[0])
+        o3d.io.write_point_cloud(f'{args.out_dir}/cloud.ply', pcd)
+        logging.info(f'Point cloud saved to {args.out_dir}')
+
+        if args.denoise_cloud:
+            logging.info('Denoising point cloud...')
+            _, ind = pcd.remove_radius_outlier(
+                nb_points=args.denoise_nb_points,
+                radius=args.denoise_radius)
+            pcd = pcd.select_by_index(ind)
+            o3d.io.write_point_cloud(f'{args.out_dir}/cloud_denoise.ply', pcd)
+
+        logging.info('Visualizing point cloud. Press ESC to exit.')
+        vis = o3d.visualization.Visualizer()
+        vis.create_window()
+        vis.add_geometry(pcd)
+        vis.get_render_option().point_size = 1.0
+        vis.get_render_option().background_color = np.array([0.5, 0.5, 0.5])
+        ctr = vis.get_view_control()
+        ctr.set_front([0, 0, -1])
+        closest = np.asarray(pcd.points)[:, 2].argmin()
+        ctr.set_lookat(np.asarray(pcd.points)[closest])
+        ctr.set_up([0, -1, 0])
+        vis.run()
+        vis.destroy_window()

cv/fast-foundationstereo/scripts/run_demo_tensorrt.py

@@ -0,0 +1,145 @@
+import os,sys
+code_dir = os.path.dirname(os.path.realpath(__file__))
+sys.path.append(f'{code_dir}/../')
+from omegaconf import OmegaConf
+from core.utils.utils import InputPadder
+import argparse, torch, logging, yaml
+import imageio
+import numpy as np
+from Utils import (
+    set_logging_format, set_seed, vis_disparity,
+    depth2xyzmap, toOpen3dCloud, o3d,
+)
+from core.foundation_stereo import TrtRunner
+import cv2
+
+
+def resolve_onnx_cfg_path(onnx_dir: str) -> str:
+  onnx_dir = os.path.normpath(onnx_dir)
+  candidates = [
+    os.path.join(onnx_dir, 'onnx.yaml'),
+    os.path.join(os.path.dirname(onnx_dir), 'onnx.yaml'),
+  ]
+  for p in candidates:
+    if os.path.exists(p):
+      return p
+  raise FileNotFoundError(
+    f"onnx.yaml not found. Looked in: {candidates}. "
+    "Please run scripts/make_onnx.py first to generate ONNX metadata."
+  )
+
+
+if __name__=="__main__":
+  code_dir = os.path.dirname(os.path.realpath(__file__))
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--onnx_dir', default=f'{code_dir}/output', type=str)
+  parser.add_argument('--left_file', default=f'{code_dir}/../assets/left.png', type=str)
+  parser.add_argument('--right_file', default=f'{code_dir}/../assets/right.png', type=str)
+  parser.add_argument('--intrinsic_file', default=f'{code_dir}/../assets/K.txt', type=str, help='camera intrinsic matrix and baseline file')
+  parser.add_argument('--out_dir', default='/home/bowen/debug/stereo_output', type=str)
+  parser.add_argument('--remove_invisible', default=1, type=int)
+  parser.add_argument('--denoise_cloud', default=1, type=int)
+  parser.add_argument('--denoise_nb_points', type=int, default=30, help='number of points to consider for radius outlier removal')
+  parser.add_argument('--denoise_radius', type=float, default=0.03, help='radius to use for outlier removal')
+  parser.add_argument('--get_pc', type=int, default=1, help='save point cloud output')
+  parser.add_argument('--zfar', type=float, default=100, help="max depth to include in point cloud")
+  args = parser.parse_args()
+
+  set_logging_format()
+  set_seed(0)
+  torch.autograd.set_grad_enabled(False)
+  os.makedirs(args.out_dir, exist_ok=True)
+
+  onnx_cfg_path = resolve_onnx_cfg_path(args.onnx_dir)
+  with open(onnx_cfg_path, 'r') as ff:
+    cfg:dict = yaml.safe_load(ff)
+  for k in args.__dict__:
+    if args.__dict__[k] is not None:
+      cfg[k] = args.__dict__[k]
+  args = OmegaConf.create(cfg)
+  logging.info(f"args:\n{args}")
+  model = TrtRunner(args, args.onnx_dir+'/feature_runner.engine', args.onnx_dir+'/post_runner.engine')
+
+  img0 = imageio.imread(args.left_file)
+  img1 = imageio.imread(args.right_file)
+  if len(img0.shape)==2:
+    img0 = np.tile(img0[...,None], (1,1,3))
+    img1 = np.tile(img1[...,None], (1,1,3))
+  img0 = img0[...,:3]
+  img1 = img1[...,:3]
+  H,W = img0.shape[:2]
+
+  fx = args.image_size[1] / img0.shape[1]
+  fy = args.image_size[0] / img0.shape[0]
+  if fx != 1 or fy != 1:
+    logging.info(f">>>>>>>>>>>>>>>WARNING: resizing image to {args.image_size}, fx: {fx}, fy: {fy}, this is not recommended. It's best to make tensorrt engine with the same image size as the input image.")
+  img0 = cv2.resize(img0, fx=fx, fy=fy, dsize=None)
+  img1 = cv2.resize(img1, fx=fx, fy=fy, dsize=None)
+  H,W = img0.shape[:2]
+  img0_ori = img0.copy()
+  img1_ori = img1.copy()
+  logging.info(f"img0: {img0.shape}")
+  imageio.imwrite(f'{args.out_dir}/left.png', img0)
+  imageio.imwrite(f'{args.out_dir}/right.png', img1)
+
+  img0 = torch.as_tensor(img0).cuda().float()[None].permute(0,3,1,2)
+  img1 = torch.as_tensor(img1).cuda().float()[None].permute(0,3,1,2)
+
+  logging.info(f"Start forward, 1st time run can be slow due to compilation")
+  disp = model.forward(img0, img1)
+  logging.info("forward done")
+  disp = disp.data.cpu().numpy().reshape(H,W).clip(0, None) * 1/fx
+
+  cmap = None
+  min_val = None
+  max_val = None
+  vis = vis_disparity(disp, min_val=min_val, max_val=max_val, cmap=cmap, color_map=cv2.COLORMAP_TURBO)
+  vis = np.concatenate([img0_ori, img1_ori, vis], axis=1)
+  imageio.imwrite(f'{args.out_dir}/disp_vis.png', vis)
+  s = 1280/vis.shape[1]
+  resized_vis = cv2.resize(vis, (int(vis.shape[1]*s), int(vis.shape[0]*s)))
+  cv2.imshow('disp', resized_vis[:,:,::-1])
+  cv2.waitKey(0)
+
+  if args.remove_invisible:
+    yy,xx = np.meshgrid(np.arange(disp.shape[0]), np.arange(disp.shape[1]), indexing='ij')
+    us_right = xx-disp
+    invalid = us_right<0
+    disp[invalid] = np.inf
+
+  if args.get_pc:
+    with open(args.intrinsic_file, 'r') as f:
+      lines = f.readlines()
+      K = np.array(list(map(float, lines[0].rstrip().split()))).astype(np.float32).reshape(3,3)
+      baseline = float(lines[1])
+    K[:2] *= np.array([fx, fy], dtype=np.float32)[:, np.newaxis]
+    depth = K[0,0]*baseline/disp
+    np.save(f'{args.out_dir}/depth_meter.npy', depth)
+    xyz_map = depth2xyzmap(depth, K)
+    pcd = toOpen3dCloud(xyz_map.reshape(-1,3), img0_ori.reshape(-1,3))
+    keep_mask = (np.asarray(pcd.points)[:,2]>0) & (np.asarray(pcd.points)[:,2]<=args.zfar)
+    keep_ids = np.arange(len(np.asarray(pcd.points)))[keep_mask]
+    pcd = pcd.select_by_index(keep_ids)
+    o3d.io.write_point_cloud(f'{args.out_dir}/cloud.ply', pcd)
+    logging.info(f"PCL saved to {args.out_dir}")
+
+    if args.denoise_cloud:
+      logging.info("[Optional step] denoise point cloud...")
+      cl, ind = pcd.remove_radius_outlier(nb_points=args.denoise_nb_points, radius=args.denoise_radius)
+      inlier_cloud = pcd.select_by_index(ind)
+      o3d.io.write_point_cloud(f'{args.out_dir}/cloud_denoise.ply', inlier_cloud)
+      pcd = inlier_cloud
+
+    logging.info("Visualizing point cloud. Press ESC to exit.")
+    vis = o3d.visualization.Visualizer()
+    vis.create_window()
+    vis.add_geometry(pcd)
+    vis.get_render_option().point_size = 1.0
+    vis.get_render_option().background_color = np.array([0.5, 0.5, 0.5])
+    ctr = vis.get_view_control()
+    ctr.set_front([0, 0, -1])
+    id = np.asarray(pcd.points)[:,2].argmin()
+    ctr.set_lookat(np.asarray(pcd.points)[id])
+    ctr.set_up([0, -1, 0])
+    vis.run()
+    vis.destroy_window()