Добавлены пропсы конвейера и стереодвижки, задействованные в прогоне
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>
This commit is contained in:
dasha_f
@@ -0,0 +1,156 @@
|
||||
cmake_minimum_required(VERSION 3.18)
|
||||
|
||||
project(fast_foundation_stereo_cpp LANGUAGES CXX CUDA)
|
||||
|
||||
set(CMAKE_CXX_STANDARD 17)
|
||||
set(CMAKE_CXX_STANDARD_REQUIRED ON)
|
||||
set(CMAKE_CUDA_STANDARD 17)
|
||||
set(CMAKE_CUDA_STANDARD_REQUIRED ON)
|
||||
|
||||
if(NOT CMAKE_BUILD_TYPE)
|
||||
set(CMAKE_BUILD_TYPE "Release" CACHE STRING
|
||||
"Build type: Debug, Release, RelWithDebInfo" FORCE)
|
||||
endif()
|
||||
|
||||
if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
|
||||
set(CMAKE_CUDA_ARCHITECTURES "80;86;89;90")
|
||||
endif()
|
||||
|
||||
add_compile_options(
|
||||
$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr>
|
||||
$<$<COMPILE_LANGUAGE:CUDA>:--diag-suppress=20015>
|
||||
$<$<COMPILE_LANGUAGE:CUDA>:--diag-suppress=20013>
|
||||
$<$<COMPILE_LANGUAGE:CUDA>:--diag-suppress=20011>
|
||||
$<$<COMPILE_LANGUAGE:CUDA>:--diag-suppress=20091>
|
||||
)
|
||||
|
||||
find_package(CUDAToolkit REQUIRED)
|
||||
|
||||
# ---- TensorRT ----
|
||||
set(TENSORRT_ROOT "/usr" CACHE PATH "TensorRT installation path")
|
||||
set(TENSORRT_HINT_DIRS
|
||||
${TENSORRT_ROOT}
|
||||
$ENV{CONDA_PREFIX}
|
||||
/usr
|
||||
/usr/local
|
||||
/usr/local/tensorrt
|
||||
${CUDAToolkit_TARGET_DIR}
|
||||
/usr/local/cuda
|
||||
/usr/local/cuda/targets/x86_64-linux
|
||||
/usr/local/cuda-13.0
|
||||
/usr/local/cuda-13.0/targets/x86_64-linux
|
||||
/usr/src/tensorrt
|
||||
)
|
||||
if(DEFINED ENV{CONDA_PREFIX})
|
||||
file(GLOB TENSORRT_PY_LIB_DIRS
|
||||
"$ENV{CONDA_PREFIX}/lib/python*/site-packages/tensorrt_libs"
|
||||
)
|
||||
file(GLOB TENSORRT_PY_INCLUDE_DIRS
|
||||
"$ENV{CONDA_PREFIX}/lib/python*/site-packages/tensorrt/include"
|
||||
)
|
||||
endif()
|
||||
|
||||
find_library(NVINFER_LIB
|
||||
NAMES nvinfer nvinfer.so.10
|
||||
HINTS
|
||||
${TENSORRT_HINT_DIRS}
|
||||
${TENSORRT_PY_LIB_DIRS}
|
||||
PATH_SUFFIXES
|
||||
lib
|
||||
lib64
|
||||
lib/x86_64-linux-gnu
|
||||
)
|
||||
if(NOT NVINFER_LIB)
|
||||
message(FATAL_ERROR "TensorRT nvinfer library not found. Set -DTENSORRT_ROOT=/path/to/tensorrt")
|
||||
endif()
|
||||
message(STATUS "TensorRT nvinfer: ${NVINFER_LIB}")
|
||||
|
||||
find_library(NVONNXPARSER_LIB
|
||||
NAMES nvonnxparser nvonnxparser.so.10
|
||||
HINTS
|
||||
${TENSORRT_HINT_DIRS}
|
||||
${TENSORRT_PY_LIB_DIRS}
|
||||
PATH_SUFFIXES
|
||||
lib
|
||||
lib64
|
||||
lib/x86_64-linux-gnu
|
||||
)
|
||||
if(NOT NVONNXPARSER_LIB)
|
||||
message(FATAL_ERROR "TensorRT nvonnxparser library not found. Set -DTENSORRT_ROOT=/path/to/tensorrt")
|
||||
endif()
|
||||
message(STATUS "TensorRT nvonnxparser: ${NVONNXPARSER_LIB}")
|
||||
|
||||
find_path(NVINFER_INCLUDE_DIR NvInfer.h
|
||||
HINTS
|
||||
${TENSORRT_HINT_DIRS}
|
||||
${TENSORRT_PY_INCLUDE_DIRS}
|
||||
PATH_SUFFIXES
|
||||
include
|
||||
include/x86_64-linux-gnu
|
||||
)
|
||||
if(NOT NVINFER_INCLUDE_DIR)
|
||||
message(FATAL_ERROR "TensorRT headers (NvInfer.h) not found. Set -DTENSORRT_ROOT=/path/to/tensorrt")
|
||||
endif()
|
||||
|
||||
find_package(OpenCV QUIET COMPONENTS core imgcodecs imgproc)
|
||||
|
||||
add_library(ffs_gwc_plugin SHARED
|
||||
src/gwc_volume_plugin.cpp
|
||||
src/depth_kernels.cu
|
||||
)
|
||||
set_target_properties(ffs_gwc_plugin PROPERTIES
|
||||
CUDA_SEPARABLE_COMPILATION ON
|
||||
CUDA_RESOLVE_DEVICE_SYMBOLS ON
|
||||
OUTPUT_NAME "ffs_gwc_plugin"
|
||||
)
|
||||
target_include_directories(ffs_gwc_plugin PUBLIC
|
||||
${CMAKE_CURRENT_SOURCE_DIR}/include
|
||||
${NVINFER_INCLUDE_DIR}
|
||||
)
|
||||
target_link_libraries(ffs_gwc_plugin PUBLIC
|
||||
${NVINFER_LIB}
|
||||
CUDA::cudart
|
||||
)
|
||||
|
||||
add_library(ffs_depth_inference STATIC
|
||||
src/ffs_depth_tensorrt.cpp
|
||||
src/ffs_depth_single_tensorrt.cpp
|
||||
src/gwc_volume_plugin.cpp
|
||||
src/depth_kernels.cu
|
||||
)
|
||||
set_target_properties(ffs_depth_inference PROPERTIES
|
||||
CUDA_SEPARABLE_COMPILATION ON
|
||||
CUDA_RESOLVE_DEVICE_SYMBOLS ON
|
||||
)
|
||||
target_include_directories(ffs_depth_inference PUBLIC
|
||||
${CMAKE_CURRENT_SOURCE_DIR}/include
|
||||
${NVINFER_INCLUDE_DIR}
|
||||
)
|
||||
target_link_libraries(ffs_depth_inference PUBLIC
|
||||
${NVINFER_LIB}
|
||||
CUDA::cudart
|
||||
)
|
||||
|
||||
if(OpenCV_FOUND)
|
||||
add_executable(ffs_depth_main app/main.cpp)
|
||||
target_link_libraries(ffs_depth_main PRIVATE
|
||||
ffs_depth_inference
|
||||
${OpenCV_LIBS}
|
||||
)
|
||||
else()
|
||||
message(WARNING "OpenCV C++ development package not found; skipping ffs_depth_main and ffs_profile_speed")
|
||||
endif()
|
||||
|
||||
add_executable(ffs_build_single_engine app/build_single_engine.cpp)
|
||||
target_link_libraries(ffs_build_single_engine PRIVATE
|
||||
ffs_depth_inference
|
||||
${NVONNXPARSER_LIB}
|
||||
)
|
||||
|
||||
if(OpenCV_FOUND)
|
||||
add_executable(ffs_profile_speed app/profile_speed.cpp)
|
||||
target_link_libraries(ffs_profile_speed PRIVATE
|
||||
ffs_depth_inference
|
||||
${OpenCV_LIBS}
|
||||
)
|
||||
endif()
|
||||
@@ -0,0 +1,249 @@
|
||||
# Fast-FoundationStereo C++ Inference
|
||||
|
||||
C++ runtime for Fast-FoundationStereo stereo depth inference on TensorRT.
|
||||
|
||||
## Run from inside the C++ Docker container
|
||||
|
||||
**All commands in this README are intended to be run from inside the container built from [`docker/dockerfile_cpp`](../docker/dockerfile_cpp).** That image carries the CUDA toolkit, TensorRT 10 runtime + ONNX parser + C++ headers, OpenCV development package, `trtexec`, and a Python environment with PyTorch / ONNX export tooling. The plain `docker/dockerfile` image does **not** ship the C++ TRT headers and will fail to build the C++ targets.
|
||||
|
||||
Environment setup:
|
||||
|
||||
```bash
|
||||
docker build --network host -t ffs -f docker/dockerfile_cpp .
|
||||
bash docker/run_container.sh
|
||||
```
|
||||
|
||||
|
||||
`run_container.sh` mounts the parent of `docker/` (the repo root) into `/workspace`, so `cd /workspace/<repo-name>/cpp` lands you in this folder. Every `cd cpp`, `cpp/build/...`, `python3 scripts/...`, and `trtexec ...` example below assumes that working directory.
|
||||
|
||||
```text
|
||||
cpp/
|
||||
+-- CMakeLists.txt
|
||||
+-- README.md
|
||||
+-- app/
|
||||
| +-- main.cpp -- ffs_depth_main: single image inference + visualization
|
||||
| +-- build_single_engine.cpp -- ffs_build_single_engine: ONNX (with FFSGWCVolume plugin) -> .engine
|
||||
| +-- profile_speed.cpp -- ffs_profile_speed: latency profiler for either route
|
||||
+-- include/
|
||||
| +-- ffs_depth_tensorrt.hpp
|
||||
| +-- ffs_depth_single_tensorrt.hpp
|
||||
| +-- ffs_gwc_plugin.hpp
|
||||
+-- src/
|
||||
+-- ffs_depth_tensorrt.cpp
|
||||
+-- ffs_depth_single_tensorrt.cpp
|
||||
+-- gwc_volume_plugin.cpp
|
||||
+-- depth_kernels.cu
|
||||
```
|
||||
|
||||
Two interchangeable inference routes are supported:
|
||||
|
||||
- `FFSSingleEngineInference` loads **one** TensorRT engine that contains the `FFSGWCVolume` plugin node.
|
||||
- `FFSDepthInference` loads **two** TensorRT engines (feature_runner + post_runner) and computes the GWC cost volume between them with a hand-written CUDA kernel.
|
||||
|
||||
`ffs_depth_main` auto-detects the route based on the engine directory contents:
|
||||
|
||||
- If the directory contains `fast_foundationstereo.engine`, it uses the single-engine plugin path.
|
||||
- Otherwise it uses the two-engine reference path.
|
||||
|
||||
Inputs and outputs are the same in both cases:
|
||||
|
||||
- Input: a stereo image pair plus an intrinsic file (`demo_data/K.txt` format: 9 floats on line 1 for the 3x3 camera matrix, one float on line 2 for the stereo baseline in meters).
|
||||
- Output: float32 disparity (input-pixel units), float32 depth in meters, and PNG visualizations.
|
||||
|
||||
## Dependencies
|
||||
|
||||
All C++ build dependencies are provided by the [`docker/dockerfile_cpp`](../docker/dockerfile_cpp) image (see [Run from inside the C++ Docker container](#run-from-inside-the-c-docker-container) for how to build and enter it):
|
||||
|
||||
- CUDA Toolkit (for `nvcc` and the CUDA runtime).
|
||||
- TensorRT 10 runtime, ONNX parser, and C++ headers.
|
||||
- OpenCV development package (image I/O and depth visualization in `app/main.cpp`).
|
||||
- `trtexec`, used by Route B to build the two-engine TensorRT engines from ONNX.
|
||||
|
||||
## Build
|
||||
|
||||
```bash
|
||||
cd cpp
|
||||
cmake -B build
|
||||
cmake --build build -j
|
||||
```
|
||||
|
||||
This produces:
|
||||
|
||||
- `build/libffs_gwc_plugin.so` -- the FFSGWCVolume plugin as a shared library (loadable by `trtexec --staticPlugins=...`).
|
||||
- `build/libffs_depth_inference.a` -- the inference static library.
|
||||
- `build/ffs_build_single_engine` -- single-engine builder.
|
||||
- `build/ffs_depth_main` -- demo / inference CLI.
|
||||
- `build/ffs_profile_speed` -- latency profiler.
|
||||
|
||||
## Route A: Single Engine with FFSGWCVolume Plugin
|
||||
|
||||
### A.1 Export ONNX with the plugin node
|
||||
|
||||
`scripts/make_plugin_onnx.py` exports one ONNX graph in which the GWC cost volume is represented by an `FFSGWCVolume` custom plugin node (resolved at engine-build time by `libffs_gwc_plugin.so`):
|
||||
|
||||
```bash
|
||||
python3 scripts/make_plugin_onnx.py \
|
||||
--model_dir weights/23-36-37/model_best_bp2_serialize.pth \
|
||||
--save_path output_plugin_onnx \
|
||||
--height 480 \
|
||||
--width 640 \
|
||||
--valid_iters 8 \
|
||||
--max_disp 192
|
||||
```
|
||||
|
||||
This writes:
|
||||
|
||||
```text
|
||||
output_plugin_onnx/
|
||||
+-- fast_foundationstereo_plugin.onnx
|
||||
+-- onnx.yaml
|
||||
```
|
||||
|
||||
### A.2 Build the single TensorRT engine
|
||||
|
||||
```bash
|
||||
cpp/build/ffs_build_single_engine \
|
||||
output_plugin_onnx/fast_foundationstereo_plugin.onnx \
|
||||
output_plugin_onnx/fast_foundationstereo.engine
|
||||
```
|
||||
|
||||
By default the engine is built with FP16 enabled. Pass `--fp32` to disable FP16. Pass `--workspace-mb N` to override the workspace (default 4096 MB).
|
||||
|
||||
After this step `output_plugin_onnx/` contains everything the runtime needs:
|
||||
|
||||
```text
|
||||
output_plugin_onnx/
|
||||
+-- fast_foundationstereo.engine
|
||||
+-- onnx.yaml
|
||||
```
|
||||
|
||||
### A.3 Run inference
|
||||
|
||||
Run from the repository root:
|
||||
|
||||
```bash
|
||||
cpp/build/ffs_depth_main \
|
||||
output_plugin_onnx \
|
||||
demo_data/left.png \
|
||||
demo_data/right.png \
|
||||
demo_data/K.txt \
|
||||
output_plugin_onnx
|
||||
```
|
||||
|
||||
The last argument is the output directory (default `ffs_output`). `ffs_depth_main` sees `fast_foundationstereo.engine` inside `output_plugin_onnx/` and uses `FFSSingleEngineInference`, which deserializes the engine and registers the `FFSGWCVolume` plugin before inference.
|
||||
|
||||
### A.4 Python alternative (build engine and run inference without the C++ apps)
|
||||
|
||||
The same plugin ONNX can be turned into an engine and executed end-to-end from Python. Only the C++ plugin shared library (`libffs_gwc_plugin.so`) is required from the C++ build; the C++ apps (`ffs_build_single_engine`, `ffs_depth_main`) are not.
|
||||
|
||||
```bash
|
||||
# 1. Export plugin ONNX
|
||||
python3 scripts/make_plugin_onnx.py \
|
||||
--model_dir weights/23-36-37/model_best_bp2_serialize.pth \
|
||||
--save_path output_plugin_onnx \
|
||||
--height 480 \
|
||||
--width 640
|
||||
|
||||
# 2. Build the C++ plugin shared library
|
||||
cmake -S cpp -B cpp/build
|
||||
cmake --build cpp/build -j
|
||||
|
||||
# 3. Build the TensorRT engine from Python
|
||||
python3 scripts/build_plugin_trt.py \
|
||||
output_plugin_onnx/fast_foundationstereo_plugin.onnx \
|
||||
output_plugin_onnx/fast_foundationstereo.engine
|
||||
|
||||
# 4. Run inference from Python
|
||||
python3 scripts/run_demo_plugin_trt.py \
|
||||
--model_dir output_plugin_onnx \
|
||||
--left_file demo_data/left.png \
|
||||
--right_file demo_data/right.png \
|
||||
--intrinsic_file demo_data/K.txt \
|
||||
--out_dir output_plugin_onnx
|
||||
```
|
||||
|
||||
Both `build_plugin_trt.py` and `run_demo_plugin_trt.py` auto-discover `libffs_gwc_plugin.so` in `cpp/build/`. Pass `--plugin_lib /path/to/libffs_gwc_plugin.so` to override the location. `build_plugin_trt.py` accepts `--fp32` and `--workspace-mb N` with the same meaning as `ffs_build_single_engine`. `run_demo_plugin_trt.py` writes the same five output files as the C++ demo (see [Outputs](#outputs)).
|
||||
|
||||
## Route B: Two TensorRT Engines (feature_runner + post_runner)
|
||||
|
||||
### B.1 Build the two engines
|
||||
|
||||
See the [Two-stage ONNX section in the top-level README](../readme.md#two-stage-onnx) for how to export `feature_runner.onnx` / `post_runner.onnx` with `scripts/make_onnx.py` and then build `feature_runner.engine` / `post_runner.engine` with `trtexec`. After running those steps your engine directory should contain:
|
||||
|
||||
```text
|
||||
output_two_onnx/
|
||||
+-- feature_runner.engine
|
||||
+-- post_runner.engine
|
||||
+-- onnx.yaml
|
||||
```
|
||||
|
||||
The two engines do **not** use the FFSGWCVolume plugin: the GWC cost volume is computed externally on GPU by `cpp/src/depth_kernels.cu` between the two engine calls, so a plain `trtexec --fp16` build with no custom plugin library is enough.
|
||||
|
||||
### B.2 Run inference
|
||||
|
||||
```bash
|
||||
cpp/build/ffs_depth_main \
|
||||
output_two_onnx \
|
||||
demo_data/left.png \
|
||||
demo_data/right.png \
|
||||
demo_data/K.txt \
|
||||
output_two_onnx
|
||||
```
|
||||
|
||||
`ffs_depth_main` sees no `fast_foundationstereo.engine` in the directory and falls back to `FFSDepthInference`, which executes the feature engine, builds the GWC volume with the CUDA kernel, then executes the post engine.
|
||||
|
||||
## Outputs
|
||||
|
||||
Both routes write the same files into the output directory:
|
||||
|
||||
- `disparity.bin` -- raw float32 disparity (input-pixel units), prefixed by int32 `[height, width]`.
|
||||
- `depth_meter.bin` -- raw float32 depth in meters, prefixed by int32 `[height, width]`.
|
||||
- `depth_meter.npy` -- NumPy float32 depth in meters, shape `[height, width]`.
|
||||
- `disp_vis.png` -- left/right/colorized-disparity side-by-side visualization.
|
||||
- `depth_vis.png` -- colorized depth visualization.
|
||||
|
||||
## Profile Speed
|
||||
|
||||
`ffs_profile_speed` benchmarks either route end-to-end on a single image pair using CUDA events for GPU-side timing and `steady_clock` for host wall time. It auto-detects the route the same way `ffs_depth_main` does (presence of `fast_foundationstereo.engine` in the engine directory selects single-engine), and `--mode` can be set explicitly:
|
||||
|
||||
```bash
|
||||
cpp/build/ffs_profile_speed <engine_dir> <left_image> <right_image> <intrinsic_file> \
|
||||
[--mode auto|two|single] [--warmup N] [--runs N] [--include-depth]
|
||||
```
|
||||
|
||||
Defaults: `--mode auto`, `--warmup 10`, `--runs 30`. With `--include-depth`, the disparity-to-depth conversion is included in the timed region (otherwise only the `infer()` call is timed).
|
||||
|
||||
### Profile the single engine
|
||||
|
||||
```bash
|
||||
cpp/build/ffs_profile_speed \
|
||||
output_plugin_onnx \
|
||||
demo_data/left.png demo_data/right.png demo_data/K.txt \
|
||||
--mode single --warmup 20 --runs 100
|
||||
```
|
||||
|
||||
### Profile the two engines
|
||||
|
||||
```bash
|
||||
cpp/build/ffs_profile_speed \
|
||||
output_two_onnx \
|
||||
demo_data/left.png demo_data/right.png demo_data/K.txt \
|
||||
--mode two --warmup 20 --runs 100
|
||||
```
|
||||
|
||||
The output looks like:
|
||||
|
||||
```text
|
||||
mode=single
|
||||
image=960x540
|
||||
model=640x480
|
||||
warmup=20 runs=100
|
||||
timed_region=infer
|
||||
gpu mean_ms=... p50_ms=... p90_ms=... min_ms=... max_ms=... std_ms=...
|
||||
host mean_ms=... p50_ms=... p90_ms=... min_ms=... max_ms=... std_ms=...
|
||||
```
|
||||
|
||||
- `gpu` is `cudaEventElapsedTime` between start/stop events on the inference stream (pure GPU work).
|
||||
- `host` is `std::chrono::steady_clock` around the same region, including the cost of `cudaEventSynchronize`. `host >= gpu` always.
|
||||
|
||||
For a side-by-side comparison, run the profiler twice with different engine directories and `--mode`, then compare the `gpu mean_ms` columns.
|
||||
@@ -0,0 +1,120 @@
|
||||
#include "ffs_gwc_plugin.hpp"
|
||||
|
||||
#include <NvInfer.h>
|
||||
#include <NvOnnxParser.h>
|
||||
|
||||
#include <filesystem>
|
||||
#include <fstream>
|
||||
#include <iostream>
|
||||
#include <memory>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
|
||||
namespace {
|
||||
|
||||
class Logger : public nvinfer1::ILogger {
|
||||
public:
|
||||
void log(Severity severity, const char* msg) noexcept override {
|
||||
if (severity <= Severity::kINFO) {
|
||||
std::cerr << "[TRT] " << msg << "\n";
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
struct TrtDestroy {
|
||||
template <typename T>
|
||||
void operator()(T* p) const {
|
||||
delete p;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
using TrtPtr = std::unique_ptr<T, TrtDestroy>;
|
||||
|
||||
void printUsage(const char* prog) {
|
||||
std::cerr
|
||||
<< "Usage: " << prog << " <plugin_onnx> <output_engine> [--fp32] [--workspace-mb N]\n"
|
||||
<< "\n"
|
||||
<< "Builds a single TensorRT engine from an ONNX graph containing the\n"
|
||||
<< "FFSGWCVolume custom plugin node.\n";
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
if (argc < 3) {
|
||||
printUsage(argv[0]);
|
||||
return 1;
|
||||
}
|
||||
|
||||
try {
|
||||
const std::filesystem::path onnx_path = argv[1];
|
||||
const std::filesystem::path engine_path = argv[2];
|
||||
bool fp16 = true;
|
||||
size_t workspace_mb = 4096;
|
||||
|
||||
for (int i = 3; i < argc; ++i) {
|
||||
const std::string arg = argv[i];
|
||||
if (arg == "--fp32") {
|
||||
fp16 = false;
|
||||
} else if (arg == "--workspace-mb" && i + 1 < argc) {
|
||||
workspace_mb = static_cast<size_t>(std::stoull(argv[++i]));
|
||||
} else {
|
||||
throw std::runtime_error("unknown argument: " + arg);
|
||||
}
|
||||
}
|
||||
|
||||
if (!std::filesystem::exists(onnx_path)) {
|
||||
throw std::runtime_error("ONNX file does not exist: " + onnx_path.string());
|
||||
}
|
||||
if (!engine_path.parent_path().empty()) {
|
||||
std::filesystem::create_directories(engine_path.parent_path());
|
||||
}
|
||||
|
||||
if (!ffs_depth::registerFFSGWCPlugin()) {
|
||||
throw std::runtime_error("failed to register FFSGWCVolume plugin");
|
||||
}
|
||||
|
||||
Logger logger;
|
||||
TrtPtr<nvinfer1::IBuilder> builder(nvinfer1::createInferBuilder(logger));
|
||||
if (!builder) throw std::runtime_error("createInferBuilder failed");
|
||||
|
||||
const auto explicit_batch =
|
||||
1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
|
||||
TrtPtr<nvinfer1::INetworkDefinition> network(builder->createNetworkV2(explicit_batch));
|
||||
if (!network) throw std::runtime_error("createNetworkV2 failed");
|
||||
|
||||
TrtPtr<nvonnxparser::IParser> parser(nvonnxparser::createParser(*network, logger));
|
||||
if (!parser) throw std::runtime_error("createParser failed");
|
||||
|
||||
if (!parser->parseFromFile(onnx_path.string().c_str(),
|
||||
static_cast<int32_t>(nvinfer1::ILogger::Severity::kWARNING))) {
|
||||
for (int32_t i = 0; i < parser->getNbErrors(); ++i) {
|
||||
std::cerr << parser->getError(i)->desc() << "\n";
|
||||
}
|
||||
throw std::runtime_error("failed to parse ONNX: " + onnx_path.string());
|
||||
}
|
||||
|
||||
TrtPtr<nvinfer1::IBuilderConfig> config(builder->createBuilderConfig());
|
||||
if (!config) throw std::runtime_error("createBuilderConfig failed");
|
||||
config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE,
|
||||
workspace_mb * 1024ULL * 1024ULL);
|
||||
if (fp16 && builder->platformHasFastFp16()) {
|
||||
config->setFlag(nvinfer1::BuilderFlag::kFP16);
|
||||
}
|
||||
|
||||
TrtPtr<nvinfer1::IHostMemory> serialized(builder->buildSerializedNetwork(*network, *config));
|
||||
if (!serialized) throw std::runtime_error("buildSerializedNetwork failed");
|
||||
|
||||
std::ofstream out(engine_path, std::ios::binary);
|
||||
if (!out) throw std::runtime_error("cannot write engine: " + engine_path.string());
|
||||
out.write(static_cast<char const*>(serialized->data()), serialized->size());
|
||||
|
||||
std::cout << "Built engine: " << engine_path << "\n";
|
||||
std::cout << "Precision: " << (fp16 ? "FP16 allowed" : "FP32") << "\n";
|
||||
} catch (const std::exception& e) {
|
||||
std::cerr << "ERROR: " << e.what() << "\n";
|
||||
return 1;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,346 @@
|
||||
/**
|
||||
* Standalone Fast-FoundationStereo TensorRT inference demo.
|
||||
*
|
||||
* Usage:
|
||||
* ./ffs_depth_main <engine_dir> <left_image> <right_image> <intrinsic_file> [output_dir]
|
||||
*
|
||||
* Outputs:
|
||||
* <output_dir>/disparity.bin int32 H, int32 W, then H*W float32 disparity
|
||||
* <output_dir>/depth_meter.bin int32 H, int32 W, then H*W float32 depth in meters
|
||||
* <output_dir>/depth_meter.npy NumPy float32 depth in meters, shape (H, W)
|
||||
* <output_dir>/disp_vis.png left/right/colorized-disparity visualization
|
||||
* <output_dir>/depth_vis.png colorized depth visualization
|
||||
*/
|
||||
|
||||
#include "ffs_depth_tensorrt.hpp"
|
||||
#include "ffs_depth_single_tensorrt.hpp"
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <opencv2/core.hpp>
|
||||
#include <opencv2/imgcodecs.hpp>
|
||||
#include <opencv2/imgproc.hpp>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstdint>
|
||||
#include <cmath>
|
||||
#include <filesystem>
|
||||
#include <fstream>
|
||||
#include <iostream>
|
||||
#include <sstream>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace {
|
||||
|
||||
void checkCuda(cudaError_t status, const char* what) {
|
||||
if (status != cudaSuccess) {
|
||||
throw std::runtime_error(std::string(what) + ": " + cudaGetErrorString(status));
|
||||
}
|
||||
}
|
||||
|
||||
void printUsage(const char* prog) {
|
||||
std::cerr
|
||||
<< "Usage: " << prog << " <engine_dir> <left_image> <right_image> <intrinsic_file> [output_dir]\n"
|
||||
<< "\n"
|
||||
<< " engine_dir : directory containing feature_runner.engine, post_runner.engine, onnx.yaml\n"
|
||||
<< " left_image : left stereo image readable by OpenCV\n"
|
||||
<< " right_image: right stereo image with the same size as left_image\n"
|
||||
<< " intrinsic_file: text file with 3x3 K on line 1 and baseline in meters on line 2\n"
|
||||
<< " output_dir : output directory (default: ffs_output)\n";
|
||||
}
|
||||
|
||||
struct CudaBuffer {
|
||||
void* ptr = nullptr;
|
||||
|
||||
~CudaBuffer() {
|
||||
if (ptr) cudaFree(ptr);
|
||||
}
|
||||
|
||||
CudaBuffer() = default;
|
||||
CudaBuffer(const CudaBuffer&) = delete;
|
||||
CudaBuffer& operator=(const CudaBuffer&) = delete;
|
||||
|
||||
void allocate(size_t bytes) {
|
||||
checkCuda(cudaMalloc(&ptr, bytes), "cudaMalloc");
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
T* as() {
|
||||
return static_cast<T*>(ptr);
|
||||
}
|
||||
};
|
||||
|
||||
struct Intrinsics {
|
||||
float k[9] = {};
|
||||
float baseline = 0.0f;
|
||||
};
|
||||
|
||||
Intrinsics loadIntrinsics(const std::string& path) {
|
||||
std::ifstream in(path);
|
||||
if (!in) {
|
||||
throw std::runtime_error("cannot open intrinsic file: " + path);
|
||||
}
|
||||
|
||||
std::string k_line;
|
||||
std::string baseline_line;
|
||||
if (!std::getline(in, k_line) || !std::getline(in, baseline_line)) {
|
||||
throw std::runtime_error("intrinsic file must contain K on line 1 and baseline on line 2");
|
||||
}
|
||||
|
||||
Intrinsics intr;
|
||||
std::istringstream k_stream(k_line);
|
||||
for (float& value : intr.k) {
|
||||
if (!(k_stream >> value)) {
|
||||
throw std::runtime_error("intrinsic file K line must contain 9 float values");
|
||||
}
|
||||
}
|
||||
|
||||
std::istringstream baseline_stream(baseline_line);
|
||||
if (!(baseline_stream >> intr.baseline)) {
|
||||
throw std::runtime_error("intrinsic file baseline line must contain one float value");
|
||||
}
|
||||
if (!(intr.k[0] > 0.0f) || !(intr.baseline > 0.0f)) {
|
||||
throw std::runtime_error("invalid focal length or baseline in intrinsic file");
|
||||
}
|
||||
|
||||
return intr;
|
||||
}
|
||||
|
||||
void saveFloatMatrix(const std::filesystem::path& path,
|
||||
const std::vector<float>& values,
|
||||
int height,
|
||||
int width) {
|
||||
std::ofstream out(path, std::ios::binary);
|
||||
if (!out) {
|
||||
throw std::runtime_error("cannot write matrix: " + path.string());
|
||||
}
|
||||
const int32_t dims[2] = {height, width};
|
||||
out.write(reinterpret_cast<const char*>(dims), sizeof(dims));
|
||||
out.write(reinterpret_cast<const char*>(values.data()),
|
||||
values.size() * sizeof(float));
|
||||
}
|
||||
|
||||
void saveNpyFloat32(const std::filesystem::path& path,
|
||||
const std::vector<float>& values,
|
||||
int height,
|
||||
int width) {
|
||||
std::ofstream out(path, std::ios::binary);
|
||||
if (!out) {
|
||||
throw std::runtime_error("cannot write npy matrix: " + path.string());
|
||||
}
|
||||
|
||||
std::ostringstream header_stream;
|
||||
header_stream << "{'descr': '<f4', 'fortran_order': False, 'shape': ("
|
||||
<< height << ", " << width << "), }";
|
||||
std::string header = header_stream.str();
|
||||
|
||||
const size_t prefix_len = 10; // magic(6) + version(2) + header_len(2)
|
||||
size_t padded_len = header.size() + 1;
|
||||
const size_t rem = (prefix_len + padded_len) % 16;
|
||||
if (rem != 0) {
|
||||
padded_len += 16 - rem;
|
||||
}
|
||||
header.append(padded_len - header.size() - 1, ' ');
|
||||
header.push_back('\n');
|
||||
|
||||
if (header.size() > 65535) {
|
||||
throw std::runtime_error("npy header too large: " + path.string());
|
||||
}
|
||||
const uint16_t header_len = static_cast<uint16_t>(header.size());
|
||||
|
||||
out.write("\x93NUMPY", 6);
|
||||
const char version[2] = {1, 0};
|
||||
out.write(version, 2);
|
||||
out.write(reinterpret_cast<const char*>(&header_len), sizeof(header_len));
|
||||
out.write(header.data(), static_cast<std::streamsize>(header.size()));
|
||||
out.write(reinterpret_cast<const char*>(values.data()),
|
||||
static_cast<std::streamsize>(values.size() * sizeof(float)));
|
||||
}
|
||||
|
||||
cv::Mat colorizeDisparity(const std::vector<float>& disparity, int height, int width) {
|
||||
std::vector<float> safe_disp(disparity.size(), 0.0f);
|
||||
std::vector<uint8_t> valid_mask(disparity.size(), 0);
|
||||
for (size_t i = 0; i < disparity.size(); ++i) {
|
||||
if (std::isfinite(disparity[i])) {
|
||||
safe_disp[i] = std::max(disparity[i], 0.0f);
|
||||
valid_mask[i] = 255;
|
||||
}
|
||||
}
|
||||
|
||||
cv::Mat disp_mat(height, width, CV_32FC1, safe_disp.data());
|
||||
cv::Mat valid(height, width, CV_8UC1, valid_mask.data());
|
||||
|
||||
if (cv::countNonZero(valid) == 0) {
|
||||
return cv::Mat::zeros(height, width, CV_8UC3);
|
||||
}
|
||||
|
||||
double min_val = 0.0;
|
||||
double max_val = 0.0;
|
||||
cv::minMaxLoc(disp_mat, &min_val, &max_val, nullptr, nullptr, valid);
|
||||
if (max_val <= min_val) {
|
||||
max_val = min_val + 1.0;
|
||||
}
|
||||
|
||||
cv::Mat disp_u8;
|
||||
disp_mat.convertTo(disp_u8, CV_8UC1, 255.0 / (max_val - min_val),
|
||||
-255.0 * min_val / (max_val - min_val));
|
||||
|
||||
cv::Mat colored;
|
||||
cv::applyColorMap(disp_u8, colored, cv::COLORMAP_TURBO);
|
||||
colored.setTo(cv::Scalar(0, 0, 0), ~valid);
|
||||
return colored;
|
||||
}
|
||||
|
||||
cv::Mat colorizeDepth(const std::vector<float>& depth, int height, int width) {
|
||||
std::vector<float> safe_depth(depth.size(), 0.0f);
|
||||
std::vector<uint8_t> valid_mask(depth.size(), 0);
|
||||
for (size_t i = 0; i < depth.size(); ++i) {
|
||||
if (std::isfinite(depth[i]) && depth[i] > 0.0f) {
|
||||
safe_depth[i] = depth[i];
|
||||
valid_mask[i] = 255;
|
||||
}
|
||||
}
|
||||
|
||||
cv::Mat depth_mat(height, width, CV_32FC1, safe_depth.data());
|
||||
cv::Mat valid(height, width, CV_8UC1, valid_mask.data());
|
||||
|
||||
if (cv::countNonZero(valid) == 0) {
|
||||
return cv::Mat::zeros(height, width, CV_8UC3);
|
||||
}
|
||||
|
||||
double min_val = 0.0;
|
||||
double max_val = 0.0;
|
||||
cv::minMaxLoc(depth_mat, &min_val, &max_val, nullptr, nullptr, valid);
|
||||
if (max_val <= min_val) {
|
||||
max_val = min_val + 1.0;
|
||||
}
|
||||
|
||||
cv::Mat depth_u8;
|
||||
depth_mat.convertTo(depth_u8, CV_8UC1, 255.0 / (max_val - min_val),
|
||||
-255.0 * min_val / (max_val - min_val));
|
||||
|
||||
cv::Mat colored;
|
||||
cv::applyColorMap(depth_u8, colored, cv::COLORMAP_TURBO);
|
||||
colored.setTo(cv::Scalar(0, 0, 0), ~valid);
|
||||
return colored;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
if (argc < 5 || argc > 6) {
|
||||
printUsage(argv[0]);
|
||||
return 1;
|
||||
}
|
||||
|
||||
try {
|
||||
const std::string engine_dir = argv[1];
|
||||
const std::string left_path = argv[2];
|
||||
const std::string right_path = argv[3];
|
||||
const std::string intrinsic_path = argv[4];
|
||||
const std::filesystem::path output_dir = (argc == 6) ? argv[5] : "ffs_output";
|
||||
const Intrinsics intr = loadIntrinsics(intrinsic_path);
|
||||
|
||||
cv::Mat left_bgr = cv::imread(left_path, cv::IMREAD_COLOR);
|
||||
cv::Mat right_bgr = cv::imread(right_path, cv::IMREAD_COLOR);
|
||||
|
||||
if (left_bgr.empty()) {
|
||||
throw std::runtime_error("cannot read left image: " + left_path);
|
||||
}
|
||||
if (right_bgr.empty()) {
|
||||
throw std::runtime_error("cannot read right image: " + right_path);
|
||||
}
|
||||
if (left_bgr.size() != right_bgr.size()) {
|
||||
throw std::runtime_error("left and right images must have identical dimensions");
|
||||
}
|
||||
if (!left_bgr.isContinuous()) left_bgr = left_bgr.clone();
|
||||
if (!right_bgr.isContinuous()) right_bgr = right_bgr.clone();
|
||||
|
||||
std::filesystem::create_directories(output_dir);
|
||||
|
||||
const int height = left_bgr.rows;
|
||||
const int width = left_bgr.cols;
|
||||
const size_t image_bytes = static_cast<size_t>(height) * width * 3 * sizeof(uint8_t);
|
||||
const size_t map_bytes = static_cast<size_t>(height) * width * sizeof(float);
|
||||
|
||||
CudaBuffer d_left;
|
||||
CudaBuffer d_right;
|
||||
CudaBuffer d_disparity;
|
||||
CudaBuffer d_depth;
|
||||
d_left.allocate(image_bytes);
|
||||
d_right.allocate(image_bytes);
|
||||
d_disparity.allocate(map_bytes);
|
||||
d_depth.allocate(map_bytes);
|
||||
|
||||
checkCuda(cudaMemcpy(d_left.ptr, left_bgr.data, image_bytes, cudaMemcpyHostToDevice),
|
||||
"cudaMemcpy left image");
|
||||
checkCuda(cudaMemcpy(d_right.ptr, right_bgr.data, image_bytes, cudaMemcpyHostToDevice),
|
||||
"cudaMemcpy right image");
|
||||
|
||||
std::cout << "Input images: " << width << "x" << height << "\n";
|
||||
std::cout << "Depth: fx=" << intr.k[0] << " baseline=" << intr.baseline << " m\n";
|
||||
std::cout << "Loading engines from: " << engine_dir << "\n";
|
||||
const bool use_single_engine =
|
||||
std::filesystem::exists(std::filesystem::path(engine_dir) / "fast_foundationstereo.engine");
|
||||
|
||||
if (use_single_engine) {
|
||||
std::cout << "Runtime: single TensorRT engine with FFSGWCVolume plugin\n";
|
||||
ffs_depth::FFSSingleEngineInference ffs(engine_dir);
|
||||
ffs.infer(d_left.as<uint8_t>(), d_right.as<uint8_t>(), height, width,
|
||||
d_disparity.as<float>());
|
||||
ffs.dispToDepth(d_disparity.as<float>(), height, width,
|
||||
intr.k[0], intr.baseline, d_depth.as<float>());
|
||||
ffs.sync();
|
||||
} else {
|
||||
std::cout << "Runtime: two TensorRT engines with external CUDA GWC\n";
|
||||
ffs_depth::FFSDepthInference ffs(engine_dir);
|
||||
ffs.infer(d_left.as<uint8_t>(), d_right.as<uint8_t>(), height, width,
|
||||
d_disparity.as<float>());
|
||||
ffs.dispToDepth(d_disparity.as<float>(), height, width,
|
||||
intr.k[0], intr.baseline, d_depth.as<float>());
|
||||
ffs.sync();
|
||||
}
|
||||
|
||||
std::vector<float> disparity(static_cast<size_t>(height) * width);
|
||||
std::vector<float> depth(static_cast<size_t>(height) * width);
|
||||
checkCuda(cudaMemcpy(disparity.data(), d_disparity.ptr, map_bytes, cudaMemcpyDeviceToHost),
|
||||
"cudaMemcpy disparity");
|
||||
checkCuda(cudaMemcpy(depth.data(), d_depth.ptr, map_bytes, cudaMemcpyDeviceToHost),
|
||||
"cudaMemcpy depth");
|
||||
|
||||
const auto disparity_path = output_dir / "disparity.bin";
|
||||
saveFloatMatrix(disparity_path, disparity, height, width);
|
||||
|
||||
const auto depth_path = output_dir / "depth_meter.bin";
|
||||
saveFloatMatrix(depth_path, depth, height, width);
|
||||
|
||||
const auto depth_npy_path = output_dir / "depth_meter.npy";
|
||||
saveNpyFloat32(depth_npy_path, depth, height, width);
|
||||
|
||||
const cv::Mat disp_color = colorizeDisparity(disparity, height, width);
|
||||
cv::Mat disp_vis;
|
||||
cv::hconcat(std::vector<cv::Mat>{left_bgr, right_bgr, disp_color}, disp_vis);
|
||||
const auto disp_vis_path = output_dir / "disp_vis.png";
|
||||
if (!cv::imwrite(disp_vis_path.string(), disp_vis)) {
|
||||
throw std::runtime_error("failed to write disparity visualization: " + disp_vis_path.string());
|
||||
}
|
||||
|
||||
const cv::Mat depth_vis = colorizeDepth(depth, height, width);
|
||||
const auto depth_vis_path = output_dir / "depth_vis.png";
|
||||
if (!cv::imwrite(depth_vis_path.string(), depth_vis)) {
|
||||
throw std::runtime_error("failed to write depth visualization: " + depth_vis_path.string());
|
||||
}
|
||||
|
||||
std::cout << "Saved: " << disparity_path << "\n";
|
||||
std::cout << "Saved: " << depth_path << "\n";
|
||||
std::cout << "Saved: " << depth_npy_path << "\n";
|
||||
std::cout << "Saved: " << disp_vis_path << "\n";
|
||||
std::cout << "Saved: " << depth_vis_path << "\n";
|
||||
} catch (const std::exception& e) {
|
||||
std::cerr << "ERROR: " << e.what() << "\n";
|
||||
return 1;
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,312 @@
|
||||
#include "ffs_depth_tensorrt.hpp"
|
||||
#include "ffs_depth_single_tensorrt.hpp"
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <opencv2/core.hpp>
|
||||
#include <opencv2/imgcodecs.hpp>
|
||||
|
||||
#include <algorithm>
|
||||
#include <chrono>
|
||||
#include <cmath>
|
||||
#include <cstdint>
|
||||
#include <filesystem>
|
||||
#include <fstream>
|
||||
#include <iostream>
|
||||
#include <numeric>
|
||||
#include <sstream>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace {
|
||||
|
||||
void checkCuda(cudaError_t status, const char* what) {
|
||||
if (status != cudaSuccess) {
|
||||
throw std::runtime_error(std::string(what) + ": " + cudaGetErrorString(status));
|
||||
}
|
||||
}
|
||||
|
||||
struct CudaBuffer {
|
||||
void* ptr = nullptr;
|
||||
~CudaBuffer() { if (ptr) cudaFree(ptr); }
|
||||
void allocate(size_t bytes) { checkCuda(cudaMalloc(&ptr, bytes), "cudaMalloc"); }
|
||||
template <typename T> T* as() { return static_cast<T*>(ptr); }
|
||||
};
|
||||
|
||||
struct Intrinsics {
|
||||
float k[9] = {};
|
||||
float baseline = 0.0f;
|
||||
};
|
||||
|
||||
struct Args {
|
||||
std::string engine_dir;
|
||||
std::string left_path;
|
||||
std::string right_path;
|
||||
std::string intrinsic_path;
|
||||
std::string mode = "auto";
|
||||
int warmup = 10;
|
||||
int runs = 30;
|
||||
bool include_depth = false;
|
||||
};
|
||||
|
||||
struct Stats {
|
||||
double mean = 0.0;
|
||||
double min = 0.0;
|
||||
double p50 = 0.0;
|
||||
double p90 = 0.0;
|
||||
double max = 0.0;
|
||||
double stddev = 0.0;
|
||||
};
|
||||
|
||||
void printUsage(const char* prog) {
|
||||
std::cerr
|
||||
<< "Usage: " << prog
|
||||
<< " <engine_dir> <left_image> <right_image> <intrinsic_file>"
|
||||
<< " [--mode auto|two|single] [--warmup N] [--runs N] [--include-depth]\n";
|
||||
}
|
||||
|
||||
Args parseArgs(int argc, char** argv) {
|
||||
if (argc < 5) {
|
||||
printUsage(argv[0]);
|
||||
throw std::runtime_error("missing required arguments");
|
||||
}
|
||||
|
||||
Args args;
|
||||
args.engine_dir = argv[1];
|
||||
args.left_path = argv[2];
|
||||
args.right_path = argv[3];
|
||||
args.intrinsic_path = argv[4];
|
||||
|
||||
for (int i = 5; i < argc; ++i) {
|
||||
const std::string key = argv[i];
|
||||
auto requireValue = [&](const char* name) -> std::string {
|
||||
if (i + 1 >= argc) {
|
||||
throw std::runtime_error(std::string("missing value for ") + name);
|
||||
}
|
||||
return argv[++i];
|
||||
};
|
||||
|
||||
if (key == "--mode") {
|
||||
args.mode = requireValue("--mode");
|
||||
} else if (key == "--warmup") {
|
||||
args.warmup = std::stoi(requireValue("--warmup"));
|
||||
} else if (key == "--runs") {
|
||||
args.runs = std::stoi(requireValue("--runs"));
|
||||
} else if (key == "--include-depth") {
|
||||
args.include_depth = true;
|
||||
} else {
|
||||
throw std::runtime_error("unknown argument: " + key);
|
||||
}
|
||||
}
|
||||
|
||||
if (args.mode != "auto" && args.mode != "two" && args.mode != "single") {
|
||||
throw std::runtime_error("--mode must be auto, two, or single");
|
||||
}
|
||||
if (args.warmup < 0 || args.runs <= 0) {
|
||||
throw std::runtime_error("--warmup must be >= 0 and --runs must be > 0");
|
||||
}
|
||||
return args;
|
||||
}
|
||||
|
||||
Intrinsics loadIntrinsics(const std::string& path) {
|
||||
std::ifstream in(path);
|
||||
if (!in) throw std::runtime_error("cannot open intrinsic file: " + path);
|
||||
std::string k_line;
|
||||
std::string baseline_line;
|
||||
if (!std::getline(in, k_line) || !std::getline(in, baseline_line)) {
|
||||
throw std::runtime_error("intrinsic file must contain K on line 1 and baseline on line 2");
|
||||
}
|
||||
|
||||
Intrinsics intr;
|
||||
std::istringstream ks(k_line);
|
||||
for (float& v : intr.k) {
|
||||
if (!(ks >> v)) throw std::runtime_error("K line must contain 9 floats");
|
||||
}
|
||||
std::istringstream bs(baseline_line);
|
||||
if (!(bs >> intr.baseline)) throw std::runtime_error("baseline line must contain one float");
|
||||
return intr;
|
||||
}
|
||||
|
||||
Stats summarize(std::vector<float> values) {
|
||||
if (values.empty()) throw std::runtime_error("cannot summarize empty timing vector");
|
||||
std::sort(values.begin(), values.end());
|
||||
const double sum = std::accumulate(values.begin(), values.end(), 0.0);
|
||||
const double mean = sum / static_cast<double>(values.size());
|
||||
double var = 0.0;
|
||||
for (float v : values) {
|
||||
const double d = static_cast<double>(v) - mean;
|
||||
var += d * d;
|
||||
}
|
||||
var /= static_cast<double>(values.size());
|
||||
|
||||
auto percentile = [&](double p) {
|
||||
const size_t idx = static_cast<size_t>(
|
||||
std::llround((values.size() - 1) * p / 100.0));
|
||||
return static_cast<double>(values[std::min(idx, values.size() - 1)]);
|
||||
};
|
||||
|
||||
Stats stats;
|
||||
stats.mean = mean;
|
||||
stats.min = values.front();
|
||||
stats.p50 = percentile(50.0);
|
||||
stats.p90 = percentile(90.0);
|
||||
stats.max = values.back();
|
||||
stats.stddev = std::sqrt(var);
|
||||
return stats;
|
||||
}
|
||||
|
||||
void printStats(const char* label, const Stats& s) {
|
||||
std::cout
|
||||
<< label
|
||||
<< " mean_ms=" << s.mean
|
||||
<< " p50_ms=" << s.p50
|
||||
<< " p90_ms=" << s.p90
|
||||
<< " min_ms=" << s.min
|
||||
<< " max_ms=" << s.max
|
||||
<< " std_ms=" << s.stddev
|
||||
<< "\n";
|
||||
}
|
||||
|
||||
template <typename Runner>
|
||||
std::vector<float> profileRunner(
|
||||
Runner& runner,
|
||||
uint8_t* d_left,
|
||||
uint8_t* d_right,
|
||||
int height,
|
||||
int width,
|
||||
float* d_disp,
|
||||
float* d_depth,
|
||||
float fx,
|
||||
float baseline,
|
||||
int warmup,
|
||||
int runs,
|
||||
bool include_depth,
|
||||
std::vector<float>& host_ms)
|
||||
{
|
||||
cudaEvent_t start = nullptr;
|
||||
cudaEvent_t stop = nullptr;
|
||||
checkCuda(cudaEventCreate(&start), "cudaEventCreate start");
|
||||
checkCuda(cudaEventCreate(&stop), "cudaEventCreate stop");
|
||||
|
||||
std::vector<float> gpu_ms;
|
||||
gpu_ms.reserve(static_cast<size_t>(runs));
|
||||
host_ms.clear();
|
||||
host_ms.reserve(static_cast<size_t>(runs));
|
||||
|
||||
for (int i = 0; i < warmup + runs; ++i) {
|
||||
const auto host_start = std::chrono::steady_clock::now();
|
||||
checkCuda(cudaEventRecord(start, runner.stream()), "cudaEventRecord start");
|
||||
runner.infer(d_left, d_right, height, width, d_disp);
|
||||
if (include_depth) {
|
||||
runner.dispToDepth(d_disp, height, width, fx, baseline, d_depth);
|
||||
}
|
||||
checkCuda(cudaEventRecord(stop, runner.stream()), "cudaEventRecord stop");
|
||||
checkCuda(cudaEventSynchronize(stop), "cudaEventSynchronize stop");
|
||||
const auto host_stop = std::chrono::steady_clock::now();
|
||||
|
||||
if (i >= warmup) {
|
||||
float elapsed = 0.0f;
|
||||
checkCuda(cudaEventElapsedTime(&elapsed, start, stop), "cudaEventElapsedTime");
|
||||
gpu_ms.push_back(elapsed);
|
||||
host_ms.push_back(static_cast<float>(
|
||||
std::chrono::duration<double, std::milli>(host_stop - host_start).count()));
|
||||
}
|
||||
}
|
||||
|
||||
cudaEventDestroy(stop);
|
||||
cudaEventDestroy(start);
|
||||
return gpu_ms;
|
||||
}
|
||||
|
||||
template <typename Runner>
|
||||
void runProfile(const Args& args,
|
||||
Runner& runner,
|
||||
uint8_t* d_left,
|
||||
uint8_t* d_right,
|
||||
int height,
|
||||
int width,
|
||||
float* d_disp,
|
||||
float* d_depth,
|
||||
const Intrinsics& intr,
|
||||
const char* mode_name) {
|
||||
std::vector<float> host_ms;
|
||||
const std::vector<float> gpu_ms = profileRunner(
|
||||
runner,
|
||||
d_left,
|
||||
d_right,
|
||||
height,
|
||||
width,
|
||||
d_disp,
|
||||
d_depth,
|
||||
intr.k[0],
|
||||
intr.baseline,
|
||||
args.warmup,
|
||||
args.runs,
|
||||
args.include_depth,
|
||||
host_ms);
|
||||
|
||||
std::cout << "mode=" << mode_name << "\n";
|
||||
std::cout << "image=" << width << "x" << height << "\n";
|
||||
std::cout << "model=" << runner.modelWidth() << "x" << runner.modelHeight() << "\n";
|
||||
std::cout << "warmup=" << args.warmup << " runs=" << args.runs << "\n";
|
||||
std::cout << "timed_region=" << (args.include_depth ? "infer+dispToDepth" : "infer") << "\n";
|
||||
printStats("gpu", summarize(gpu_ms));
|
||||
printStats("host", summarize(host_ms));
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
try {
|
||||
const Args args = parseArgs(argc, argv);
|
||||
const Intrinsics intr = loadIntrinsics(args.intrinsic_path);
|
||||
|
||||
cv::Mat left = cv::imread(args.left_path, cv::IMREAD_COLOR);
|
||||
cv::Mat right = cv::imread(args.right_path, cv::IMREAD_COLOR);
|
||||
if (left.empty()) throw std::runtime_error("cannot read left image: " + args.left_path);
|
||||
if (right.empty()) throw std::runtime_error("cannot read right image: " + args.right_path);
|
||||
if (left.size() != right.size()) throw std::runtime_error("left/right size mismatch");
|
||||
if (!left.isContinuous()) left = left.clone();
|
||||
if (!right.isContinuous()) right = right.clone();
|
||||
|
||||
const int height = left.rows;
|
||||
const int width = left.cols;
|
||||
const size_t image_bytes = static_cast<size_t>(height) * width * 3;
|
||||
const size_t map_bytes = static_cast<size_t>(height) * width * sizeof(float);
|
||||
|
||||
CudaBuffer d_left;
|
||||
CudaBuffer d_right;
|
||||
CudaBuffer d_disp;
|
||||
CudaBuffer d_depth;
|
||||
d_left.allocate(image_bytes);
|
||||
d_right.allocate(image_bytes);
|
||||
d_disp.allocate(map_bytes);
|
||||
d_depth.allocate(map_bytes);
|
||||
checkCuda(cudaMemcpy(d_left.ptr, left.data, image_bytes, cudaMemcpyHostToDevice), "copy left");
|
||||
checkCuda(cudaMemcpy(d_right.ptr, right.data, image_bytes, cudaMemcpyHostToDevice), "copy right");
|
||||
|
||||
std::string mode = args.mode;
|
||||
if (mode == "auto") {
|
||||
mode = std::filesystem::exists(
|
||||
std::filesystem::path(args.engine_dir) / "fast_foundationstereo.engine")
|
||||
? "single"
|
||||
: "two";
|
||||
}
|
||||
|
||||
if (mode == "single") {
|
||||
ffs_depth::FFSSingleEngineInference runner(args.engine_dir);
|
||||
runProfile(args, runner, d_left.as<uint8_t>(), d_right.as<uint8_t>(),
|
||||
height, width, d_disp.as<float>(), d_depth.as<float>(),
|
||||
intr, "single");
|
||||
} else {
|
||||
ffs_depth::FFSDepthInference runner(args.engine_dir);
|
||||
runProfile(args, runner, d_left.as<uint8_t>(), d_right.as<uint8_t>(),
|
||||
height, width, d_disp.as<float>(), d_depth.as<float>(),
|
||||
intr, "two");
|
||||
}
|
||||
} catch (const std::exception& e) {
|
||||
std::cerr << "ERROR: " << e.what() << "\n";
|
||||
return 1;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,80 @@
|
||||
#pragma once
|
||||
|
||||
#include <NvInfer.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cstdint>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
|
||||
#include "ffs_depth_tensorrt.hpp"
|
||||
|
||||
namespace ffs_depth {
|
||||
|
||||
/**
|
||||
* Single TensorRT engine inference path.
|
||||
*
|
||||
* Expected engine directory layout:
|
||||
* - fast_foundationstereo.engine
|
||||
* - onnx.yaml
|
||||
*
|
||||
* The engine is built from the plugin ONNX export path where the GWC cost
|
||||
* volume is represented by an FFSGWCVolume TensorRT plugin node.
|
||||
*/
|
||||
class FFSSingleEngineInference {
|
||||
public:
|
||||
explicit FFSSingleEngineInference(const std::string& engine_dir);
|
||||
~FFSSingleEngineInference();
|
||||
|
||||
FFSSingleEngineInference(const FFSSingleEngineInference&) = delete;
|
||||
FFSSingleEngineInference& operator=(const FFSSingleEngineInference&) = delete;
|
||||
|
||||
void infer(const uint8_t* d_left_rgb, const uint8_t* d_right_rgb,
|
||||
int input_h, int input_w,
|
||||
float* d_disp_out);
|
||||
|
||||
void dispToDepth(const float* d_disp,
|
||||
int height, int width,
|
||||
float fx, float baseline_m,
|
||||
float* d_depth_out);
|
||||
|
||||
void inferDepth(const uint8_t* d_left_rgb, const uint8_t* d_right_rgb,
|
||||
int input_h, int input_w,
|
||||
float fx, float baseline_m,
|
||||
float* d_depth_out);
|
||||
|
||||
void sync();
|
||||
|
||||
int modelHeight() const { return config_.image_height; }
|
||||
int modelWidth() const { return config_.image_width; }
|
||||
int maxDisp() const { return config_.max_disp; }
|
||||
int cvGroup() const { return config_.cv_group; }
|
||||
const FFSDepthInference::Config& config() const { return config_; }
|
||||
cudaStream_t stream() const { return stream_; }
|
||||
|
||||
private:
|
||||
void loadConfig(const std::string& config_path);
|
||||
void loadEngine(const std::string& path);
|
||||
void allocateBuffers();
|
||||
void freeDeviceBuffers();
|
||||
void preprocessRGBGPU(const uint8_t* d_rgb, int src_h, int src_w, float* d_output);
|
||||
|
||||
FFSDepthInference::Config config_;
|
||||
|
||||
std::unique_ptr<nvinfer1::IRuntime> runtime_;
|
||||
std::unique_ptr<nvinfer1::ICudaEngine> engine_;
|
||||
std::unique_ptr<nvinfer1::IExecutionContext> context_;
|
||||
|
||||
cudaStream_t stream_ = nullptr;
|
||||
|
||||
float* d_left_ = nullptr;
|
||||
float* d_right_ = nullptr;
|
||||
float* d_disp_ = nullptr;
|
||||
float* d_disp_cropped_ = nullptr;
|
||||
float* d_disp_for_depth_ = nullptr;
|
||||
|
||||
int64_t depth_alloc_pixels_ = 0;
|
||||
int scaled_w_ = 0;
|
||||
int scaled_h_ = 0;
|
||||
};
|
||||
|
||||
} // namespace ffs_depth
|
||||
@@ -0,0 +1,190 @@
|
||||
#pragma once
|
||||
|
||||
#include <NvInfer.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cstdint>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace ffs_depth {
|
||||
|
||||
/**
|
||||
* FoundationStereo TensorRT depth inference.
|
||||
*
|
||||
* Uses the two-engine architecture (feature_runner + post_runner).
|
||||
* The GWC (Group-wise Correlation) volume is computed on GPU between the two engines.
|
||||
*
|
||||
* Expected engine directory layout:
|
||||
* - feature_runner.engine
|
||||
* - post_runner.engine
|
||||
* - onnx.yaml (contains image_size, max_disp, cv_group, valid_iters)
|
||||
*/
|
||||
class FFSDepthInference {
|
||||
public:
|
||||
struct Config {
|
||||
int image_height = 480; // Model input height (must be divisible by 32)
|
||||
int image_width = 864; // Model input width (must be divisible by 32)
|
||||
int max_disp = 192; // Maximum disparity search range in pixels
|
||||
int cv_group = 8; // Number of groups for group-wise correlation (GWC) volume
|
||||
int valid_iters = 8; // Number of GRU refinement iterations during inference
|
||||
bool normalize = true; // Whether to L2-normalize features in GWC correlation
|
||||
};
|
||||
|
||||
/**
|
||||
* Load engines and allocate all GPU buffers.
|
||||
* @param engine_dir Directory containing the two .engine files and onnx.yaml.
|
||||
*/
|
||||
explicit FFSDepthInference(const std::string& engine_dir);
|
||||
~FFSDepthInference();
|
||||
|
||||
FFSDepthInference(const FFSDepthInference&) = delete;
|
||||
FFSDepthInference& operator=(const FFSDepthInference&) = delete;
|
||||
|
||||
/**
|
||||
* Run stereo disparity inference entirely on the GPU.
|
||||
*
|
||||
* This method is asynchronous: all work (preprocessing, TensorRT
|
||||
* inference, postprocessing) is enqueued on an internal CUDA stream
|
||||
* and the call returns immediately. Call sync() to block until the
|
||||
* output buffer is safe to read.
|
||||
*
|
||||
* Input images are expected in HWC uint8 BGR (OpenCV) format on device memory.
|
||||
* If (input_h, input_w) differs from the model size, the images are
|
||||
* uniformly (aspect-preserving) bilinearly resized down to fit the model
|
||||
* input and the remaining right/bottom strip is filled with replicate
|
||||
* padding. After inference the disparity is cropped back to the scaled
|
||||
* region and nearest-neighbour upsampled to (input_h, input_w), with
|
||||
* a horizontal scale correction so the values are in input-pixel units.
|
||||
*
|
||||
* @param d_left_rgb Left image on GPU (input_h x input_w x 3, uint8 BGR)
|
||||
* @param d_right_rgb Right image on GPU (input_h x input_w x 3, uint8 BGR)
|
||||
* @param input_h Height of input images (must be > 0)
|
||||
* @param input_w Width of input images (must be > 0)
|
||||
* @param d_disp_out Output disparity on GPU (input_h x input_w, float32,
|
||||
* in input-pixel units, clamped to >= 0).
|
||||
*/
|
||||
void infer(const uint8_t* d_left_rgb, const uint8_t* d_right_rgb,
|
||||
int input_h, int input_w,
|
||||
float* d_disp_out);
|
||||
|
||||
/**
|
||||
* Convert an input-resolution disparity map to float32 depth in meters.
|
||||
*
|
||||
* This method is asynchronous: work is enqueued on the internal CUDA
|
||||
* stream and the call returns immediately. Call sync() before reading.
|
||||
*
|
||||
* Conversion semantics match scripts/run_demo.py (see inferDepth docstring
|
||||
* for the +inf / off-image-correspondence handling).
|
||||
*
|
||||
* @param d_disp Input disparity on GPU (height x width, float32).
|
||||
* Must be expressed in INPUT-image pixel units (which
|
||||
* is what infer() produces).
|
||||
* @param height Disparity / depth height
|
||||
* @param width Disparity / depth width
|
||||
* @param fx Focal length in pixels at INPUT resolution
|
||||
* @param baseline_m Stereo baseline in meters
|
||||
* @param d_depth_out Output depth on GPU (height x width, float32, meters)
|
||||
*/
|
||||
void dispToDepth(const float* d_disp,
|
||||
int height, int width,
|
||||
float fx, float baseline_m,
|
||||
float* d_depth_out);
|
||||
|
||||
/**
|
||||
* Run stereo inference and convert disparity to float32 depth (meters) in one call.
|
||||
*
|
||||
* This method is asynchronous: all work is enqueued on an internal
|
||||
* CUDA stream and the call returns immediately. Call sync() to block
|
||||
* until d_depth_out is safe to read.
|
||||
*
|
||||
* The conversion matches scripts/run_demo.py:
|
||||
* depth_m = fx * baseline_m / disparity
|
||||
* where:
|
||||
* - disparity is in INPUT-image pixel units (infer() already rescales it),
|
||||
* so `fx` must come from the INPUT-resolution intrinsics (i.e. K[0,0]
|
||||
* of the unscaled camera matrix, not the model-resolution intrinsics).
|
||||
* - pixels with x - disparity < 0 (right-image correspondence off-image)
|
||||
* are marked invalid: disparity is replaced with +inf, yielding depth 0.
|
||||
* - disparity == 0 yields depth = +inf (consumers should mask non-finite
|
||||
* values before using the depth map for downstream geometry).
|
||||
*
|
||||
* @param d_left_rgb Left image on GPU (input_h x input_w x 3, uint8 BGR/RGB)
|
||||
* @param d_right_rgb Right image on GPU (input_h x input_w x 3, uint8 BGR/RGB)
|
||||
* @param input_h Height of input images
|
||||
* @param input_w Width of input images
|
||||
* @param fx Focal length in pixels at INPUT resolution
|
||||
* @param baseline_m Stereo baseline in meters
|
||||
* @param d_depth_out Output depth on GPU (input_h x input_w, float32, meters)
|
||||
*/
|
||||
void inferDepth(const uint8_t* d_left_rgb, const uint8_t* d_right_rgb,
|
||||
int input_h, int input_w,
|
||||
float fx, float baseline_m,
|
||||
float* d_depth_out);
|
||||
|
||||
/** Block until all async work on the internal CUDA stream has completed. */
|
||||
void sync();
|
||||
|
||||
int modelHeight() const { return config_.image_height; }
|
||||
int modelWidth() const { return config_.image_width; }
|
||||
int maxDisp() const { return config_.max_disp; }
|
||||
int cvGroup() const { return config_.cv_group; }
|
||||
const Config& config() const { return config_; }
|
||||
cudaStream_t stream() const { return stream_; }
|
||||
|
||||
private:
|
||||
void loadConfig(const std::string& config_path);
|
||||
void loadEngine(const std::string& path,
|
||||
std::unique_ptr<nvinfer1::ICudaEngine>& engine,
|
||||
std::unique_ptr<nvinfer1::IExecutionContext>& context);
|
||||
void allocateBuffers();
|
||||
void allocateFeatureBuffers();
|
||||
void allocatePostBuffers();
|
||||
void freeDeviceBuffers(); // safe to call repeatedly; nulls every pointer
|
||||
void preprocessRGBGPU(const uint8_t* d_rgb, int src_h, int src_w, float* d_output);
|
||||
|
||||
void runFeatureRunner();
|
||||
void buildGWCVolume();
|
||||
void runPostRunner();
|
||||
|
||||
Config config_;
|
||||
|
||||
std::unique_ptr<nvinfer1::IRuntime> runtime_;
|
||||
std::unique_ptr<nvinfer1::ICudaEngine> feature_engine_;
|
||||
std::unique_ptr<nvinfer1::IExecutionContext> feature_context_;
|
||||
std::unique_ptr<nvinfer1::ICudaEngine> post_engine_;
|
||||
std::unique_ptr<nvinfer1::IExecutionContext> post_context_;
|
||||
|
||||
cudaStream_t stream_ = nullptr;
|
||||
|
||||
float* d_left_ = nullptr;
|
||||
float* d_right_ = nullptr;
|
||||
|
||||
float* d_feat_left_04_ = nullptr;
|
||||
float* d_feat_left_08_ = nullptr;
|
||||
float* d_feat_left_16_ = nullptr;
|
||||
float* d_feat_left_32_ = nullptr;
|
||||
float* d_feat_right_04_ = nullptr;
|
||||
float* d_stem_2x_ = nullptr;
|
||||
|
||||
std::vector<int> feat_04_dims_;
|
||||
std::vector<int> feat_08_dims_;
|
||||
std::vector<int> feat_16_dims_;
|
||||
std::vector<int> feat_32_dims_;
|
||||
std::vector<int> stem_2x_dims_;
|
||||
|
||||
float* d_gwc_volume_ = nullptr;
|
||||
int gwc_disp_levels_ = 0;
|
||||
|
||||
float* d_disp_ = nullptr;
|
||||
float* d_disp_cropped_ = nullptr;
|
||||
float* d_disp_for_depth_ = nullptr;
|
||||
int64_t depth_alloc_pixels_ = 0;
|
||||
|
||||
int scaled_w_ = 0; // Uniform-scaled width (before padding)
|
||||
int scaled_h_ = 0; // Uniform-scaled height (before padding)
|
||||
|
||||
bool gwc_fp16_ = false; // GWC volume tensor is FP16 (vs. FP32)
|
||||
};
|
||||
|
||||
} // namespace ffs_depth
|
||||
@@ -0,0 +1,12 @@
|
||||
#pragma once
|
||||
|
||||
namespace ffs_depth {
|
||||
|
||||
// Registers the FFSGWCVolume TensorRT plugin creator in the global registry.
|
||||
// Safe to call more than once.
|
||||
bool registerFFSGWCPlugin();
|
||||
|
||||
} // namespace ffs_depth
|
||||
|
||||
// C ABI wrapper for loading/registering the plugin from Python via ctypes.
|
||||
extern "C" bool ffs_register_gwc_plugin();
|
||||
@@ -0,0 +1,345 @@
|
||||
/**
|
||||
* CUDA kernels for the Fast-FoundationStereo TensorRT depth inference pipeline.
|
||||
*
|
||||
* Active kernels:
|
||||
* 1. GWC volume (mixed precision) -- FP32 features -> FP16 correlation volume
|
||||
* 2. HWC uint8 -> CHW float -- exact-match path (no resize, with zero-padding)
|
||||
* 3. Uniform resize + pad -- aspect-ratio-preserving bilinear resize with
|
||||
* border-replicate padding to model dimensions
|
||||
* 4. Disparity crop -- removes padding from model-resolution disparity
|
||||
* 5. Disparity upsample -- nearest-neighbor upsample with scale correction
|
||||
* 6. Disparity clamp -- clamps to minimum value
|
||||
* 7. Disparity to depth -- depth_m = fx * baseline_m / disparity (float32, meters)
|
||||
*/
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <device_launch_parameters.h>
|
||||
#include <math_constants.h>
|
||||
#include <cstdint>
|
||||
|
||||
namespace ffs_depth {
|
||||
namespace cuda {
|
||||
|
||||
// =========================================================================
|
||||
// 1. GWC Volume
|
||||
// =========================================================================
|
||||
template <typename T>
|
||||
__device__ __forceinline__ float toFloat(T v) {
|
||||
return static_cast<float>(v);
|
||||
}
|
||||
|
||||
template <>
|
||||
__device__ __forceinline__ float toFloat<__half>(__half v) {
|
||||
return __half2float(v);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__device__ __forceinline__ T fromFloat(float v) {
|
||||
return static_cast<T>(v);
|
||||
}
|
||||
|
||||
template <>
|
||||
__device__ __forceinline__ __half fromFloat<__half>(float v) {
|
||||
return __float2half(v);
|
||||
}
|
||||
|
||||
template <typename InputT, typename OutputT>
|
||||
__global__ void buildGWCVolumeKernel(
|
||||
const InputT* __restrict__ feat_left,
|
||||
const InputT* __restrict__ feat_right,
|
||||
OutputT* __restrict__ gwc_volume,
|
||||
int B, int C, int H, int W,
|
||||
int max_disp, int num_groups, bool normalize)
|
||||
{
|
||||
const int w = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int h = blockIdx.y * blockDim.y + threadIdx.y;
|
||||
const int dgb = blockIdx.z;
|
||||
|
||||
if (w >= W || h >= H) return;
|
||||
|
||||
const int d = dgb % max_disp;
|
||||
const int g = (dgb / max_disp) % num_groups;
|
||||
const int b = dgb / (max_disp * num_groups);
|
||||
if (b >= B) return;
|
||||
|
||||
const int K = C / num_groups;
|
||||
const int w_right = w - d;
|
||||
|
||||
const int out_idx = ((b * num_groups + g) * max_disp + d) * H * W + h * W + w;
|
||||
|
||||
if (w_right < 0) { gwc_volume[out_idx] = fromFloat<OutputT>(0.0f); return; }
|
||||
|
||||
float dot = 0.f, nl = 0.f, nr = 0.f;
|
||||
const int left_base = (b * C + g * K) * H * W + h * W + w;
|
||||
const int right_base = (b * C + g * K) * H * W + h * W + w_right;
|
||||
const int stride = H * W;
|
||||
|
||||
for (int k = 0; k < K; ++k) {
|
||||
float l = toFloat<InputT>(feat_left [left_base + k * stride]);
|
||||
float r = toFloat<InputT>(feat_right[right_base + k * stride]);
|
||||
dot += l * r; nl += l * l; nr += r * r;
|
||||
}
|
||||
if (normalize)
|
||||
gwc_volume[out_idx] = fromFloat<OutputT>(dot / (sqrtf(nl) * sqrtf(nr) + 1e-5f));
|
||||
else
|
||||
gwc_volume[out_idx] = fromFloat<OutputT>(dot);
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// 2. Preprocess RGB HWC uint8 -> CHW float (exact-match path, no resize)
|
||||
// Used when input resolution matches model resolution exactly.
|
||||
// =========================================================================
|
||||
__global__ void preprocessRGBToCHWKernel(
|
||||
const uint8_t* __restrict__ d_rgb_hwc,
|
||||
float* __restrict__ d_chw,
|
||||
int src_h, int src_w,
|
||||
int dst_h, int dst_w)
|
||||
{
|
||||
const int x = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int y = blockIdx.y * blockDim.y + threadIdx.y;
|
||||
if (x >= dst_w || y >= dst_h) return;
|
||||
|
||||
const int hw = dst_h * dst_w;
|
||||
const int idx = y * dst_w + x;
|
||||
|
||||
if (x < src_w && y < src_h) {
|
||||
const int si = (y * src_w + x) * 3;
|
||||
d_chw[idx] = static_cast<float>(d_rgb_hwc[si + 2]); // R
|
||||
d_chw[hw + idx] = static_cast<float>(d_rgb_hwc[si + 1]); // G
|
||||
d_chw[2 * hw + idx] = static_cast<float>(d_rgb_hwc[si]); // B
|
||||
} else {
|
||||
d_chw[idx] = 0.f; d_chw[hw + idx] = 0.f; d_chw[2 * hw + idx] = 0.f;
|
||||
}
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// 3. Uniform resize + border-replicate padding (aspect-ratio preserving)
|
||||
// Pixels in [0, scaled_w) x [0, scaled_h) are bilinear-sampled from src.
|
||||
// Pixels in the padding region replicate the nearest edge pixel.
|
||||
// =========================================================================
|
||||
__global__ void resizeUniformAndPadKernel(
|
||||
const uint8_t* __restrict__ d_rgb_hwc,
|
||||
float* __restrict__ d_chw,
|
||||
int src_h, int src_w,
|
||||
int scaled_h, int scaled_w,
|
||||
int dst_h, int dst_w)
|
||||
{
|
||||
const int x = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int y = blockIdx.y * blockDim.y + threadIdx.y;
|
||||
if (x >= dst_w || y >= dst_h) return;
|
||||
|
||||
const int hw = dst_h * dst_w;
|
||||
const int idx = y * dst_w + x;
|
||||
|
||||
const int cx = min(x, scaled_w - 1);
|
||||
const int cy = min(y, scaled_h - 1);
|
||||
|
||||
const float sx = static_cast<float>(src_w) / scaled_w;
|
||||
const float sy = static_cast<float>(src_h) / scaled_h;
|
||||
|
||||
const float fx = (cx + 0.5f) * sx - 0.5f;
|
||||
const float fy = (cy + 0.5f) * sy - 0.5f;
|
||||
|
||||
const int x0 = max(0, min(__float2int_rd(fx), src_w - 1));
|
||||
const int y0 = max(0, min(__float2int_rd(fy), src_h - 1));
|
||||
const int x1 = min(x0 + 1, src_w - 1);
|
||||
const int y1 = min(y0 + 1, src_h - 1);
|
||||
|
||||
const float wx = fx - floorf(fx);
|
||||
const float wy = fy - floorf(fy);
|
||||
|
||||
for (int c = 0; c < 3; ++c) {
|
||||
const int sc = 2 - c; // RGB channel reorder
|
||||
float v00 = static_cast<float>(d_rgb_hwc[(y0 * src_w + x0) * 3 + sc]);
|
||||
float v01 = static_cast<float>(d_rgb_hwc[(y0 * src_w + x1) * 3 + sc]);
|
||||
float v10 = static_cast<float>(d_rgb_hwc[(y1 * src_w + x0) * 3 + sc]);
|
||||
float v11 = static_cast<float>(d_rgb_hwc[(y1 * src_w + x1) * 3 + sc]);
|
||||
float val = (1.f - wy) * ((1.f - wx) * v00 + wx * v01)
|
||||
+ wy * ((1.f - wx) * v10 + wx * v11);
|
||||
d_chw[c * hw + idx] = val;
|
||||
}
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// 4. Crop disparity (remove border-replicate padding)
|
||||
// Extracts the valid (scaled_w x scaled_h) region from the
|
||||
// model-resolution (model_w x model_h) disparity output.
|
||||
// =========================================================================
|
||||
__global__ void cropDisparityKernel(
|
||||
const float* __restrict__ src, float* __restrict__ dst,
|
||||
int src_w, int dst_h, int dst_w)
|
||||
{
|
||||
const int x = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int y = blockIdx.y * blockDim.y + threadIdx.y;
|
||||
if (x >= dst_w || y >= dst_h) return;
|
||||
dst[y * dst_w + x] = src[y * src_w + x];
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// 5. Upsample float disparity (nearest neighbor) with scale correction
|
||||
// Upsamples from cropped model resolution to input resolution.
|
||||
// disp_scale converts disparity from scaled-pixel units to input-pixel units.
|
||||
// =========================================================================
|
||||
__global__ void upsampleDisparityKernel(
|
||||
const float* __restrict__ src, float* __restrict__ dst,
|
||||
int src_w, int src_h,
|
||||
int dst_w, int dst_h,
|
||||
float disp_scale)
|
||||
{
|
||||
const int x = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int y = blockIdx.y * blockDim.y + threadIdx.y;
|
||||
if (x >= dst_w || y >= dst_h) return;
|
||||
|
||||
const float sx = static_cast<float>(src_w) / dst_w;
|
||||
const float sy = static_cast<float>(src_h) / dst_h;
|
||||
|
||||
const float fx = (x + 0.5f) * sx - 0.5f;
|
||||
const float fy = (y + 0.5f) * sy - 0.5f;
|
||||
|
||||
const int x0 = max(0, min(__float2int_rn(fx), src_w - 1));
|
||||
const int y0 = max(0, min(__float2int_rn(fy), src_h - 1));
|
||||
float val = src[y0 * src_w + x0];
|
||||
|
||||
dst[y * dst_w + x] = val * disp_scale;
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// 6. Clamp disparity to minimum value (removes negative/zero artifacts)
|
||||
// =========================================================================
|
||||
__global__ void clampDisparityKernel(
|
||||
float* __restrict__ disp, int count, float min_val)
|
||||
{
|
||||
const int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx >= count) return;
|
||||
disp[idx] = fmaxf(disp[idx], min_val);
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// 7. Disparity to depth: depth_m = fx * baseline_m / disparity (float32 meters)
|
||||
//
|
||||
// Matches scripts/run_demo.py behaviour:
|
||||
// - disparity is assumed already clamped to >= 0 (see clampDisparityKernel)
|
||||
// - pixels where the right-image correspondence would fall off-image
|
||||
// (x - disp < 0) are marked invalid by setting disparity to +inf,
|
||||
// which yields depth = 0 after division.
|
||||
// - disp == 0 yields depth = +inf (matches numpy's float divide-by-zero).
|
||||
// - no clip on the upper end; output is float32 meters.
|
||||
// =========================================================================
|
||||
__global__ void dispToDepthKernel(
|
||||
const float* __restrict__ disp,
|
||||
float* __restrict__ depth_m,
|
||||
int height, int width, float fx, float baseline_m)
|
||||
{
|
||||
const int x = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int y = blockIdx.y * blockDim.y + threadIdx.y;
|
||||
if (x >= width || y >= height) return;
|
||||
const int i = y * width + x;
|
||||
|
||||
float d = disp[i];
|
||||
if (static_cast<float>(x) - d < 0.0f) {
|
||||
d = CUDART_INF_F; // remove_invisible: drop pixel via inf -> depth 0
|
||||
}
|
||||
depth_m[i] = fx * baseline_m / d;
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Host wrappers (extern "C" for linkage)
|
||||
// =========================================================================
|
||||
extern "C" {
|
||||
|
||||
void ffsCudaBuildGWCVolumeFloat(
|
||||
const float* d_fl, const float* d_fr, float* d_gwc,
|
||||
int B, int C, int H, int W, int max_disp, int ngroups, bool normalize, cudaStream_t s)
|
||||
{
|
||||
dim3 blk(16, 16);
|
||||
dim3 grd((W + 15) / 16, (H + 15) / 16, B * ngroups * max_disp);
|
||||
buildGWCVolumeKernel<float, float><<<grd, blk, 0, s>>>(d_fl, d_fr, d_gwc, B, C, H, W, max_disp, ngroups, normalize);
|
||||
}
|
||||
|
||||
void ffsCudaBuildGWCVolumeHalf(
|
||||
const __half* d_fl, const __half* d_fr, __half* d_gwc,
|
||||
int B, int C, int H, int W, int max_disp, int ngroups, bool normalize, cudaStream_t s)
|
||||
{
|
||||
dim3 blk(16, 16);
|
||||
dim3 grd((W + 15) / 16, (H + 15) / 16, B * ngroups * max_disp);
|
||||
buildGWCVolumeKernel<__half, __half><<<grd, blk, 0, s>>>(d_fl, d_fr, d_gwc, B, C, H, W, max_disp, ngroups, normalize);
|
||||
}
|
||||
|
||||
void ffsCudaBuildGWCVolumeHalfToFloat(
|
||||
const __half* d_fl, const __half* d_fr, float* d_gwc,
|
||||
int B, int C, int H, int W, int max_disp, int ngroups, bool normalize, cudaStream_t s)
|
||||
{
|
||||
dim3 blk(16, 16);
|
||||
dim3 grd((W + 15) / 16, (H + 15) / 16, B * ngroups * max_disp);
|
||||
buildGWCVolumeKernel<__half, float><<<grd, blk, 0, s>>>(d_fl, d_fr, d_gwc, B, C, H, W, max_disp, ngroups, normalize);
|
||||
}
|
||||
|
||||
void ffsCudaBuildGWCVolumeMixed(
|
||||
const float* d_fl, const float* d_fr, __half* d_gwc,
|
||||
int B, int C, int H, int W, int max_disp, int ngroups, bool normalize, cudaStream_t s)
|
||||
{
|
||||
dim3 blk(16, 16);
|
||||
dim3 grd((W + 15) / 16, (H + 15) / 16, B * ngroups * max_disp);
|
||||
buildGWCVolumeKernel<float, __half><<<grd, blk, 0, s>>>(d_fl, d_fr, d_gwc, B, C, H, W, max_disp, ngroups, normalize);
|
||||
}
|
||||
|
||||
void ffsCudaPreprocessRGBToCHW(
|
||||
const uint8_t* d_rgb, float* d_chw,
|
||||
int src_h, int src_w, int dst_h, int dst_w, cudaStream_t s)
|
||||
{
|
||||
dim3 blk(32, 16);
|
||||
dim3 grd((dst_w + 31) / 32, (dst_h + 15) / 16);
|
||||
preprocessRGBToCHWKernel<<<grd, blk, 0, s>>>(d_rgb, d_chw, src_h, src_w, dst_h, dst_w);
|
||||
}
|
||||
|
||||
void ffsCudaResizeUniformAndPad(
|
||||
const uint8_t* d_rgb, float* d_chw,
|
||||
int src_h, int src_w, int scaled_h, int scaled_w,
|
||||
int dst_h, int dst_w, cudaStream_t s)
|
||||
{
|
||||
dim3 blk(32, 16);
|
||||
dim3 grd((dst_w + 31) / 32, (dst_h + 15) / 16);
|
||||
resizeUniformAndPadKernel<<<grd, blk, 0, s>>>(d_rgb, d_chw, src_h, src_w, scaled_h, scaled_w, dst_h, dst_w);
|
||||
}
|
||||
|
||||
void ffsCudaCropDisparity(
|
||||
const float* d_src, float* d_dst,
|
||||
int src_h, int src_w, int dst_h, int dst_w, cudaStream_t s)
|
||||
{
|
||||
dim3 blk(32, 16);
|
||||
dim3 grd((dst_w + 31) / 32, (dst_h + 15) / 16);
|
||||
cropDisparityKernel<<<grd, blk, 0, s>>>(d_src, d_dst, src_w, dst_h, dst_w);
|
||||
}
|
||||
|
||||
void ffsCudaUpsampleDisparity(
|
||||
const float* d_src, float* d_dst,
|
||||
int src_w, int src_h, int dst_w, int dst_h,
|
||||
float disp_scale, cudaStream_t s)
|
||||
{
|
||||
dim3 blk(32, 16);
|
||||
dim3 grd((dst_w + 31) / 32, (dst_h + 15) / 16);
|
||||
upsampleDisparityKernel<<<grd, blk, 0, s>>>(d_src, d_dst, src_w, src_h, dst_w, dst_h, disp_scale);
|
||||
}
|
||||
|
||||
void ffsCudaClampDisparity(
|
||||
float* d_disp, int count, float min_val, cudaStream_t s)
|
||||
{
|
||||
int threads = 256;
|
||||
int blocks = (count + threads - 1) / threads;
|
||||
clampDisparityKernel<<<blocks, threads, 0, s>>>(d_disp, count, min_val);
|
||||
}
|
||||
|
||||
void ffsCudaDispToDepth(
|
||||
const float* d_disp, float* d_depth_m,
|
||||
int height, int width, float fx, float baseline_m, cudaStream_t s)
|
||||
{
|
||||
dim3 blk(32, 16);
|
||||
dim3 grd((width + 31) / 32, (height + 15) / 16);
|
||||
dispToDepthKernel<<<grd, blk, 0, s>>>(d_disp, d_depth_m, height, width, fx, baseline_m);
|
||||
}
|
||||
|
||||
} // extern "C"
|
||||
|
||||
} // namespace cuda
|
||||
} // namespace ffs_depth
|
||||
@@ -0,0 +1,392 @@
|
||||
#include "ffs_depth_single_tensorrt.hpp"
|
||||
|
||||
#include "ffs_gwc_plugin.hpp"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cctype>
|
||||
#include <cmath>
|
||||
#include <cstdint>
|
||||
#include <filesystem>
|
||||
#include <fstream>
|
||||
#include <iostream>
|
||||
#include <sstream>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace ffs_depth {
|
||||
|
||||
namespace cuda {
|
||||
extern "C" {
|
||||
|
||||
void ffsCudaPreprocessRGBToCHW(
|
||||
const uint8_t* d_rgb, float* d_chw,
|
||||
int src_h, int src_w, int dst_h, int dst_w, cudaStream_t s);
|
||||
|
||||
void ffsCudaResizeUniformAndPad(
|
||||
const uint8_t* d_rgb, float* d_chw,
|
||||
int src_h, int src_w, int scaled_h, int scaled_w,
|
||||
int dst_h, int dst_w, cudaStream_t s);
|
||||
|
||||
void ffsCudaCropDisparity(
|
||||
const float* d_src, float* d_dst,
|
||||
int src_h, int src_w, int dst_h, int dst_w, cudaStream_t s);
|
||||
|
||||
void ffsCudaUpsampleDisparity(
|
||||
const float* d_src, float* d_dst,
|
||||
int src_w, int src_h, int dst_w, int dst_h,
|
||||
float disp_scale, cudaStream_t s);
|
||||
|
||||
void ffsCudaClampDisparity(
|
||||
float* d_disp, int count, float min_val, cudaStream_t s);
|
||||
|
||||
void ffsCudaDispToDepth(
|
||||
const float* d_disp, float* d_depth_m,
|
||||
int height, int width, float fx, float baseline_m, cudaStream_t s);
|
||||
|
||||
} // extern "C"
|
||||
} // namespace cuda
|
||||
|
||||
namespace {
|
||||
|
||||
class TrtLogger : public nvinfer1::ILogger {
|
||||
public:
|
||||
void log(Severity severity, const char* msg) noexcept override {
|
||||
if (severity <= Severity::kWARNING) {
|
||||
std::cerr << "[TRT] " << msg << std::endl;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
TrtLogger g_trt_logger;
|
||||
|
||||
size_t elementSize(nvinfer1::DataType dt) {
|
||||
switch (dt) {
|
||||
case nvinfer1::DataType::kFLOAT: return 4;
|
||||
case nvinfer1::DataType::kHALF: return 2;
|
||||
case nvinfer1::DataType::kINT8: return 1;
|
||||
case nvinfer1::DataType::kINT32: return 4;
|
||||
default: return 4;
|
||||
}
|
||||
}
|
||||
|
||||
void cudaMallocChecked(void** ptr, size_t bytes, const char* what) {
|
||||
cudaError_t err = cudaMalloc(ptr, bytes);
|
||||
if (err != cudaSuccess) {
|
||||
throw std::runtime_error(std::string("[FFS single] cudaMalloc failed for ") +
|
||||
what + " (" + std::to_string(bytes) + " bytes): " +
|
||||
cudaGetErrorString(err));
|
||||
}
|
||||
}
|
||||
|
||||
bool hasTensor(nvinfer1::ICudaEngine& engine, const char* name) {
|
||||
for (int32_t i = 0; i < engine.getNbIOTensors(); ++i) {
|
||||
if (std::string(engine.getIOTensorName(i)) == name) return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
std::string findSingleYamlConfig(const std::string& engine_dir) {
|
||||
namespace fs = std::filesystem;
|
||||
|
||||
std::vector<fs::path> yaml_files;
|
||||
std::error_code ec;
|
||||
fs::directory_iterator it(engine_dir, ec);
|
||||
if (ec) {
|
||||
throw std::runtime_error("[FFS single] Cannot list engine directory: " +
|
||||
engine_dir + " (" + ec.message() + ")");
|
||||
}
|
||||
|
||||
for (const auto& entry : it) {
|
||||
if (ec) {
|
||||
throw std::runtime_error("[FFS single] Cannot list engine directory: " +
|
||||
engine_dir + " (" + ec.message() + ")");
|
||||
}
|
||||
if (!entry.is_regular_file()) continue;
|
||||
|
||||
std::string ext = entry.path().extension().string();
|
||||
std::transform(ext.begin(), ext.end(), ext.begin(), [](unsigned char c) {
|
||||
return static_cast<char>(std::tolower(c));
|
||||
});
|
||||
if (ext == ".yaml" || ext == ".yml") {
|
||||
yaml_files.push_back(entry.path());
|
||||
}
|
||||
}
|
||||
|
||||
if (yaml_files.empty()) {
|
||||
throw std::runtime_error("[FFS single] No YAML config found in: " + engine_dir);
|
||||
}
|
||||
if (yaml_files.size() > 1) {
|
||||
std::ostringstream oss;
|
||||
oss << "[FFS single] Expected exactly one YAML config in " << engine_dir
|
||||
<< ", found " << yaml_files.size() << ":";
|
||||
for (const auto& path : yaml_files) {
|
||||
oss << " " << path.string();
|
||||
}
|
||||
throw std::runtime_error(oss.str());
|
||||
}
|
||||
|
||||
return yaml_files.front().string();
|
||||
}
|
||||
|
||||
std::string trim(std::string s) {
|
||||
const char* ws = " \t\r\n";
|
||||
const size_t start = s.find_first_not_of(ws);
|
||||
if (start == std::string::npos) return "";
|
||||
const size_t end = s.find_last_not_of(ws);
|
||||
return s.substr(start, end - start + 1);
|
||||
}
|
||||
|
||||
std::vector<int> parseInts(std::string s) {
|
||||
for (char& c : s) {
|
||||
if (!(c >= '0' && c <= '9')) c = ' ';
|
||||
}
|
||||
std::stringstream ss(s);
|
||||
std::vector<int> values;
|
||||
int v = 0;
|
||||
while (ss >> v) values.push_back(v);
|
||||
return values;
|
||||
}
|
||||
|
||||
bool parseBool(std::string s) {
|
||||
s = trim(s);
|
||||
std::transform(s.begin(), s.end(), s.begin(), [](unsigned char c) {
|
||||
return static_cast<char>(std::tolower(c));
|
||||
});
|
||||
return s == "true" || s == "1" || s == "yes" || s == "on";
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
FFSSingleEngineInference::FFSSingleEngineInference(const std::string& engine_dir) {
|
||||
cudaError_t err = cudaStreamCreate(&stream_);
|
||||
if (err != cudaSuccess) {
|
||||
throw std::runtime_error(std::string("[FFS single] cudaStreamCreate failed: ") +
|
||||
cudaGetErrorString(err));
|
||||
}
|
||||
try {
|
||||
loadConfig(findSingleYamlConfig(engine_dir));
|
||||
if (!registerFFSGWCPlugin()) {
|
||||
throw std::runtime_error("[FFS single] failed to register FFSGWCVolume plugin");
|
||||
}
|
||||
runtime_.reset(nvinfer1::createInferRuntime(g_trt_logger));
|
||||
if (!runtime_) {
|
||||
throw std::runtime_error("[FFS single] createInferRuntime returned null");
|
||||
}
|
||||
loadEngine(engine_dir + "/fast_foundationstereo.engine");
|
||||
allocateBuffers();
|
||||
} catch (...) {
|
||||
freeDeviceBuffers();
|
||||
if (stream_) {
|
||||
cudaStreamDestroy(stream_);
|
||||
stream_ = nullptr;
|
||||
}
|
||||
throw;
|
||||
}
|
||||
}
|
||||
|
||||
FFSSingleEngineInference::~FFSSingleEngineInference() {
|
||||
freeDeviceBuffers();
|
||||
if (stream_) {
|
||||
cudaStreamDestroy(stream_);
|
||||
stream_ = nullptr;
|
||||
}
|
||||
}
|
||||
|
||||
void FFSSingleEngineInference::freeDeviceBuffers() {
|
||||
auto free = [](void*& p) { if (p) { cudaFree(p); p = nullptr; } };
|
||||
free(reinterpret_cast<void*&>(d_left_));
|
||||
free(reinterpret_cast<void*&>(d_right_));
|
||||
free(reinterpret_cast<void*&>(d_disp_));
|
||||
free(reinterpret_cast<void*&>(d_disp_cropped_));
|
||||
free(reinterpret_cast<void*&>(d_disp_for_depth_));
|
||||
depth_alloc_pixels_ = 0;
|
||||
}
|
||||
|
||||
void FFSSingleEngineInference::loadConfig(const std::string& path) {
|
||||
std::ifstream f(path);
|
||||
if (!f.good()) throw std::runtime_error("[FFS single] Cannot open config: " + path);
|
||||
|
||||
std::string line;
|
||||
int image_index = -1;
|
||||
while (std::getline(f, line)) {
|
||||
const size_t comment = line.find('#');
|
||||
if (comment != std::string::npos) line = line.substr(0, comment);
|
||||
line = trim(line);
|
||||
if (line.empty()) continue;
|
||||
|
||||
if (image_index >= 0) {
|
||||
if (line.rfind("-", 0) == 0) {
|
||||
const auto values = parseInts(line);
|
||||
if (!values.empty()) {
|
||||
if (image_index == 0) config_.image_height = values[0];
|
||||
if (image_index == 1) config_.image_width = values[0];
|
||||
++image_index;
|
||||
if (image_index >= 2) image_index = -1;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
image_index = -1;
|
||||
}
|
||||
|
||||
const size_t colon = line.find(':');
|
||||
if (colon == std::string::npos) continue;
|
||||
const std::string key = trim(line.substr(0, colon));
|
||||
const std::string value = trim(line.substr(colon + 1));
|
||||
|
||||
if (key == "image_size") {
|
||||
const auto values = parseInts(value);
|
||||
if (values.size() >= 2) {
|
||||
config_.image_height = values[0];
|
||||
config_.image_width = values[1];
|
||||
} else {
|
||||
image_index = 0;
|
||||
}
|
||||
} else if (key == "max_disp") {
|
||||
const auto values = parseInts(value);
|
||||
if (!values.empty()) config_.max_disp = values[0];
|
||||
} else if (key == "cv_group") {
|
||||
const auto values = parseInts(value);
|
||||
if (!values.empty()) config_.cv_group = values[0];
|
||||
} else if (key == "valid_iters") {
|
||||
const auto values = parseInts(value);
|
||||
if (!values.empty()) config_.valid_iters = values[0];
|
||||
} else if (key == "normalize") {
|
||||
config_.normalize = parseBool(value);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void FFSSingleEngineInference::loadEngine(const std::string& path) {
|
||||
std::ifstream f(path, std::ios::binary);
|
||||
if (!f.good()) throw std::runtime_error("[FFS single] Cannot open engine: " + path);
|
||||
|
||||
f.seekg(0, std::ios::end);
|
||||
size_t sz = f.tellg();
|
||||
f.seekg(0, std::ios::beg);
|
||||
std::vector<char> buf(sz);
|
||||
f.read(buf.data(), sz);
|
||||
|
||||
engine_.reset(runtime_->deserializeCudaEngine(buf.data(), sz));
|
||||
if (!engine_) throw std::runtime_error("[FFS single] Deserialize failed: " + path);
|
||||
|
||||
context_.reset(engine_->createExecutionContext());
|
||||
if (!context_) throw std::runtime_error("[FFS single] Context creation failed: " + path);
|
||||
|
||||
if (!hasTensor(*engine_, "left") || !hasTensor(*engine_, "right") || !hasTensor(*engine_, "disp")) {
|
||||
throw std::runtime_error("[FFS single] Engine must expose tensors named left, right, and disp");
|
||||
}
|
||||
}
|
||||
|
||||
void FFSSingleEngineInference::allocateBuffers() {
|
||||
const size_t H = static_cast<size_t>(config_.image_height);
|
||||
const size_t W = static_cast<size_t>(config_.image_width);
|
||||
const size_t input_bytes = 3 * H * W * sizeof(float);
|
||||
|
||||
cudaMallocChecked(reinterpret_cast<void**>(&d_left_), input_bytes, "d_left_");
|
||||
cudaMallocChecked(reinterpret_cast<void**>(&d_right_), input_bytes, "d_right_");
|
||||
|
||||
const size_t disp_bytes = H * W * elementSize(engine_->getTensorDataType("disp"));
|
||||
cudaMallocChecked(reinterpret_cast<void**>(&d_disp_), disp_bytes, "d_disp_");
|
||||
cudaMallocChecked(reinterpret_cast<void**>(&d_disp_cropped_),
|
||||
H * W * sizeof(float), "d_disp_cropped_");
|
||||
|
||||
context_->setTensorAddress("left", d_left_);
|
||||
context_->setTensorAddress("right", d_right_);
|
||||
context_->setTensorAddress("disp", d_disp_);
|
||||
}
|
||||
|
||||
void FFSSingleEngineInference::preprocessRGBGPU(
|
||||
const uint8_t* d_rgb, int src_h, int src_w, float* d_output) {
|
||||
const int mH = config_.image_height;
|
||||
const int mW = config_.image_width;
|
||||
if (src_h == mH && src_w == mW) {
|
||||
scaled_h_ = mH;
|
||||
scaled_w_ = mW;
|
||||
cuda::ffsCudaPreprocessRGBToCHW(d_rgb, d_output, src_h, src_w, mH, mW, stream_);
|
||||
} else {
|
||||
const float scale = std::min(static_cast<float>(mW) / src_w,
|
||||
static_cast<float>(mH) / src_h);
|
||||
scaled_w_ = std::max(1, static_cast<int>(std::round(src_w * scale)));
|
||||
scaled_h_ = std::max(1, static_cast<int>(std::round(src_h * scale)));
|
||||
cuda::ffsCudaResizeUniformAndPad(
|
||||
d_rgb, d_output, src_h, src_w, scaled_h_, scaled_w_, mH, mW, stream_);
|
||||
}
|
||||
}
|
||||
|
||||
void FFSSingleEngineInference::infer(
|
||||
const uint8_t* d_left_rgb, const uint8_t* d_right_rgb,
|
||||
int input_h, int input_w,
|
||||
float* d_disp_out) {
|
||||
if (!d_left_rgb || !d_right_rgb || !d_disp_out) {
|
||||
throw std::runtime_error("[FFS single] infer: null device pointer");
|
||||
}
|
||||
if (input_h <= 0 || input_w <= 0) {
|
||||
throw std::runtime_error("[FFS single] infer: input dimensions must be positive");
|
||||
}
|
||||
|
||||
const int mH = config_.image_height;
|
||||
const int mW = config_.image_width;
|
||||
const bool needs_resize = (input_h != mH || input_w != mW);
|
||||
|
||||
preprocessRGBGPU(d_left_rgb, input_h, input_w, d_left_);
|
||||
preprocessRGBGPU(d_right_rgb, input_h, input_w, d_right_);
|
||||
|
||||
if (!needs_resize) {
|
||||
context_->setTensorAddress("disp", d_disp_out);
|
||||
}
|
||||
if (!context_->enqueueV3(stream_)) {
|
||||
throw std::runtime_error("[FFS single] enqueue failed");
|
||||
}
|
||||
|
||||
if (needs_resize) {
|
||||
cuda::ffsCudaClampDisparity(d_disp_, mH * mW, 0.0f, stream_);
|
||||
cuda::ffsCudaCropDisparity(d_disp_, d_disp_cropped_,
|
||||
mH, mW, scaled_h_, scaled_w_, stream_);
|
||||
const float disp_scale = static_cast<float>(input_w) / scaled_w_;
|
||||
cuda::ffsCudaUpsampleDisparity(d_disp_cropped_, d_disp_out,
|
||||
scaled_w_, scaled_h_, input_w, input_h,
|
||||
disp_scale, stream_);
|
||||
} else {
|
||||
cuda::ffsCudaClampDisparity(d_disp_out, mH * mW, 0.0f, stream_);
|
||||
context_->setTensorAddress("disp", d_disp_);
|
||||
}
|
||||
}
|
||||
|
||||
void FFSSingleEngineInference::dispToDepth(
|
||||
const float* d_disp,
|
||||
int height, int width,
|
||||
float fx, float baseline_m,
|
||||
float* d_depth_out) {
|
||||
cuda::ffsCudaDispToDepth(d_disp, d_depth_out,
|
||||
height, width, fx, baseline_m, stream_);
|
||||
}
|
||||
|
||||
void FFSSingleEngineInference::inferDepth(
|
||||
const uint8_t* d_left_rgb, const uint8_t* d_right_rgb,
|
||||
int input_h, int input_w,
|
||||
float fx, float baseline_m,
|
||||
float* d_depth_out) {
|
||||
if (input_h <= 0 || input_w <= 0) {
|
||||
throw std::runtime_error("[FFS single] inferDepth: input dimensions must be positive");
|
||||
}
|
||||
const size_t num_pixels = static_cast<size_t>(input_h) * static_cast<size_t>(input_w);
|
||||
if (static_cast<int64_t>(num_pixels) > depth_alloc_pixels_) {
|
||||
if (d_disp_for_depth_) {
|
||||
cudaFree(d_disp_for_depth_);
|
||||
d_disp_for_depth_ = nullptr;
|
||||
}
|
||||
cudaMallocChecked(reinterpret_cast<void**>(&d_disp_for_depth_),
|
||||
num_pixels * sizeof(float), "d_disp_for_depth_");
|
||||
depth_alloc_pixels_ = static_cast<int64_t>(num_pixels);
|
||||
}
|
||||
infer(d_left_rgb, d_right_rgb, input_h, input_w, d_disp_for_depth_);
|
||||
cuda::ffsCudaDispToDepth(d_disp_for_depth_, d_depth_out,
|
||||
input_h, input_w, fx, baseline_m, stream_);
|
||||
}
|
||||
|
||||
void FFSSingleEngineInference::sync() {
|
||||
cudaStreamSynchronize(stream_);
|
||||
}
|
||||
|
||||
} // namespace ffs_depth
|
||||
@@ -0,0 +1,488 @@
|
||||
#include "ffs_depth_tensorrt.hpp"
|
||||
|
||||
#include <cuda_fp16.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cctype>
|
||||
#include <cmath>
|
||||
#include <cstdint>
|
||||
#include <fstream>
|
||||
#include <iostream>
|
||||
#include <sstream>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace ffs_depth {
|
||||
|
||||
// Forward declarations of CUDA kernel wrappers (defined in depth_kernels.cu)
|
||||
namespace cuda {
|
||||
extern "C" {
|
||||
|
||||
void ffsCudaBuildGWCVolumeMixed(
|
||||
const float* d_fl, const float* d_fr, __half* d_gwc,
|
||||
int B, int C, int H, int W, int max_disp, int ngroups, bool normalize, cudaStream_t s);
|
||||
|
||||
void ffsCudaPreprocessRGBToCHW(
|
||||
const uint8_t* d_rgb, float* d_chw,
|
||||
int src_h, int src_w, int dst_h, int dst_w, cudaStream_t s);
|
||||
|
||||
void ffsCudaResizeUniformAndPad(
|
||||
const uint8_t* d_rgb, float* d_chw,
|
||||
int src_h, int src_w, int scaled_h, int scaled_w,
|
||||
int dst_h, int dst_w, cudaStream_t s);
|
||||
|
||||
void ffsCudaCropDisparity(
|
||||
const float* d_src, float* d_dst,
|
||||
int src_h, int src_w, int dst_h, int dst_w, cudaStream_t s);
|
||||
|
||||
void ffsCudaUpsampleDisparity(
|
||||
const float* d_src, float* d_dst,
|
||||
int src_w, int src_h, int dst_w, int dst_h,
|
||||
float disp_scale, cudaStream_t s);
|
||||
|
||||
void ffsCudaClampDisparity(
|
||||
float* d_disp, int count, float min_val, cudaStream_t s);
|
||||
|
||||
void ffsCudaDispToDepth(
|
||||
const float* d_disp, float* d_depth_m,
|
||||
int height, int width, float fx, float baseline_m, cudaStream_t s);
|
||||
|
||||
} // extern "C"
|
||||
} // namespace cuda
|
||||
|
||||
namespace {
|
||||
|
||||
class TrtLogger : public nvinfer1::ILogger {
|
||||
public:
|
||||
void log(Severity severity, const char* msg) noexcept override {
|
||||
if (severity <= Severity::kWARNING)
|
||||
std::cerr << "[TRT] " << msg << std::endl;
|
||||
}
|
||||
};
|
||||
|
||||
TrtLogger g_trt_logger;
|
||||
|
||||
size_t elementSize(nvinfer1::DataType dt) {
|
||||
switch (dt) {
|
||||
case nvinfer1::DataType::kFLOAT: return 4;
|
||||
case nvinfer1::DataType::kHALF: return 2;
|
||||
case nvinfer1::DataType::kINT8: return 1;
|
||||
case nvinfer1::DataType::kINT32: return 4;
|
||||
default: return 4;
|
||||
}
|
||||
}
|
||||
|
||||
void cudaMallocChecked(void** ptr, size_t bytes, const char* what) {
|
||||
cudaError_t err = cudaMalloc(ptr, bytes);
|
||||
if (err != cudaSuccess) {
|
||||
throw std::runtime_error(std::string("[FFS] cudaMalloc failed for ") + what +
|
||||
" (" + std::to_string(bytes) + " bytes): " +
|
||||
cudaGetErrorString(err));
|
||||
}
|
||||
}
|
||||
|
||||
std::string trim(std::string s) {
|
||||
const char* ws = " \t\r\n";
|
||||
const size_t start = s.find_first_not_of(ws);
|
||||
if (start == std::string::npos) return "";
|
||||
const size_t end = s.find_last_not_of(ws);
|
||||
return s.substr(start, end - start + 1);
|
||||
}
|
||||
|
||||
std::vector<int> parseInts(std::string s) {
|
||||
for (char& c : s) {
|
||||
if (!(c >= '0' && c <= '9')) c = ' ';
|
||||
}
|
||||
std::stringstream ss(s);
|
||||
std::vector<int> values;
|
||||
int v = 0;
|
||||
while (ss >> v) values.push_back(v);
|
||||
return values;
|
||||
}
|
||||
|
||||
bool parseBool(std::string s) {
|
||||
s = trim(s);
|
||||
std::transform(s.begin(), s.end(), s.begin(), [](unsigned char c) {
|
||||
return static_cast<char>(std::tolower(c));
|
||||
});
|
||||
return s == "true" || s == "1" || s == "yes" || s == "on";
|
||||
}
|
||||
|
||||
} // anonymous namespace
|
||||
|
||||
// =========================================================================
|
||||
// Construction / destruction
|
||||
// =========================================================================
|
||||
|
||||
FFSDepthInference::FFSDepthInference(const std::string& engine_dir) {
|
||||
cudaError_t err = cudaStreamCreate(&stream_);
|
||||
if (err != cudaSuccess) {
|
||||
throw std::runtime_error(std::string("[FFS] cudaStreamCreate failed: ") +
|
||||
cudaGetErrorString(err));
|
||||
}
|
||||
try {
|
||||
loadConfig(engine_dir + "/onnx.yaml");
|
||||
runtime_.reset(nvinfer1::createInferRuntime(g_trt_logger));
|
||||
if (!runtime_) {
|
||||
throw std::runtime_error("[FFS] createInferRuntime returned null");
|
||||
}
|
||||
loadEngine(engine_dir + "/feature_runner.engine", feature_engine_, feature_context_);
|
||||
loadEngine(engine_dir + "/post_runner.engine", post_engine_, post_context_);
|
||||
allocateBuffers();
|
||||
} catch (...) {
|
||||
// Best-effort cleanup of anything allocated so far. Buffers may be partially
|
||||
// populated; freeDeviceBuffers() is a no-op for null pointers.
|
||||
freeDeviceBuffers();
|
||||
if (stream_) {
|
||||
cudaStreamDestroy(stream_);
|
||||
stream_ = nullptr;
|
||||
}
|
||||
throw;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
FFSDepthInference::~FFSDepthInference() {
|
||||
freeDeviceBuffers();
|
||||
if (stream_) {
|
||||
cudaStreamDestroy(stream_);
|
||||
stream_ = nullptr;
|
||||
}
|
||||
}
|
||||
|
||||
void FFSDepthInference::freeDeviceBuffers() {
|
||||
auto free = [](void*& p) { if (p) { cudaFree(p); p = nullptr; } };
|
||||
|
||||
free(reinterpret_cast<void*&>(d_left_));
|
||||
free(reinterpret_cast<void*&>(d_right_));
|
||||
free(reinterpret_cast<void*&>(d_feat_left_04_));
|
||||
free(reinterpret_cast<void*&>(d_feat_left_08_));
|
||||
free(reinterpret_cast<void*&>(d_feat_left_16_));
|
||||
free(reinterpret_cast<void*&>(d_feat_left_32_));
|
||||
free(reinterpret_cast<void*&>(d_feat_right_04_));
|
||||
free(reinterpret_cast<void*&>(d_stem_2x_));
|
||||
free(reinterpret_cast<void*&>(d_gwc_volume_));
|
||||
free(reinterpret_cast<void*&>(d_disp_));
|
||||
free(reinterpret_cast<void*&>(d_disp_cropped_));
|
||||
free(reinterpret_cast<void*&>(d_disp_for_depth_));
|
||||
depth_alloc_pixels_ = 0;
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Config / engine loading
|
||||
// =========================================================================
|
||||
|
||||
void FFSDepthInference::loadConfig(const std::string& path) {
|
||||
std::ifstream f(path);
|
||||
if (!f.good()) throw std::runtime_error("[FFS] Cannot open config: " + path);
|
||||
|
||||
std::string line;
|
||||
int image_index = -1;
|
||||
while (std::getline(f, line)) {
|
||||
const size_t comment = line.find('#');
|
||||
if (comment != std::string::npos) line = line.substr(0, comment);
|
||||
line = trim(line);
|
||||
if (line.empty()) continue;
|
||||
|
||||
if (image_index >= 0) {
|
||||
if (line.rfind("-", 0) == 0) {
|
||||
const auto values = parseInts(line);
|
||||
if (!values.empty()) {
|
||||
if (image_index == 0) config_.image_height = values[0];
|
||||
if (image_index == 1) config_.image_width = values[0];
|
||||
++image_index;
|
||||
if (image_index >= 2) image_index = -1;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
image_index = -1;
|
||||
}
|
||||
|
||||
const size_t colon = line.find(':');
|
||||
if (colon == std::string::npos) continue;
|
||||
const std::string key = trim(line.substr(0, colon));
|
||||
const std::string value = trim(line.substr(colon + 1));
|
||||
|
||||
if (key == "image_size") {
|
||||
const auto values = parseInts(value);
|
||||
if (values.size() >= 2) {
|
||||
config_.image_height = values[0];
|
||||
config_.image_width = values[1];
|
||||
} else {
|
||||
image_index = 0;
|
||||
}
|
||||
} else if (key == "max_disp") {
|
||||
const auto values = parseInts(value);
|
||||
if (!values.empty()) config_.max_disp = values[0];
|
||||
} else if (key == "cv_group") {
|
||||
const auto values = parseInts(value);
|
||||
if (!values.empty()) config_.cv_group = values[0];
|
||||
} else if (key == "valid_iters") {
|
||||
const auto values = parseInts(value);
|
||||
if (!values.empty()) config_.valid_iters = values[0];
|
||||
} else if (key == "normalize") {
|
||||
config_.normalize = parseBool(value);
|
||||
}
|
||||
}
|
||||
|
||||
gwc_disp_levels_ = config_.max_disp / 4;
|
||||
}
|
||||
|
||||
void FFSDepthInference::loadEngine(
|
||||
const std::string& path,
|
||||
std::unique_ptr<nvinfer1::ICudaEngine>& engine,
|
||||
std::unique_ptr<nvinfer1::IExecutionContext>& ctx)
|
||||
{
|
||||
std::ifstream f(path, std::ios::binary);
|
||||
if (!f.good()) throw std::runtime_error("[FFS] Cannot open engine: " + path);
|
||||
|
||||
f.seekg(0, std::ios::end);
|
||||
size_t sz = f.tellg();
|
||||
f.seekg(0, std::ios::beg);
|
||||
std::vector<char> buf(sz);
|
||||
f.read(buf.data(), sz);
|
||||
|
||||
engine.reset(runtime_->deserializeCudaEngine(buf.data(), sz));
|
||||
if (!engine) throw std::runtime_error("[FFS] Deserialize failed: " + path);
|
||||
|
||||
ctx.reset(engine->createExecutionContext());
|
||||
if (!ctx) throw std::runtime_error("[FFS] Context creation failed: " + path);
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Buffer allocation
|
||||
// =========================================================================
|
||||
|
||||
void FFSDepthInference::allocateBuffers() {
|
||||
const size_t H = static_cast<size_t>(config_.image_height);
|
||||
const size_t W = static_cast<size_t>(config_.image_width);
|
||||
const size_t input_bytes = 3 * H * W * sizeof(float);
|
||||
|
||||
cudaMallocChecked(reinterpret_cast<void**>(&d_left_), input_bytes, "d_left_");
|
||||
cudaMallocChecked(reinterpret_cast<void**>(&d_right_), input_bytes, "d_right_");
|
||||
|
||||
allocateFeatureBuffers();
|
||||
allocatePostBuffers();
|
||||
}
|
||||
|
||||
void FFSDepthInference::allocateFeatureBuffers() {
|
||||
auto getDims = [&](const char* name) {
|
||||
auto d = feature_engine_->getTensorShape(name);
|
||||
std::vector<int> v(d.nbDims);
|
||||
for (int i = 0; i < d.nbDims; ++i) v[i] = d.d[i];
|
||||
return v;
|
||||
};
|
||||
|
||||
auto allocTensor = [&](const char* name, const std::vector<int>& dims) -> float* {
|
||||
auto dt = feature_engine_->getTensorDataType(name);
|
||||
size_t bytes = elementSize(dt);
|
||||
for (int d : dims) {
|
||||
if (d <= 0) {
|
||||
throw std::runtime_error(std::string("[FFS] tensor '") + name +
|
||||
"' has non-positive dimension");
|
||||
}
|
||||
bytes *= static_cast<size_t>(d);
|
||||
}
|
||||
void* p = nullptr;
|
||||
cudaMallocChecked(&p, bytes, name);
|
||||
return static_cast<float*>(p);
|
||||
};
|
||||
|
||||
feat_04_dims_ = getDims("features_left_04");
|
||||
feat_08_dims_ = getDims("features_left_08");
|
||||
feat_16_dims_ = getDims("features_left_16");
|
||||
feat_32_dims_ = getDims("features_left_32");
|
||||
stem_2x_dims_ = getDims("stem_2x");
|
||||
|
||||
d_feat_left_04_ = allocTensor("features_left_04", feat_04_dims_);
|
||||
d_feat_left_08_ = allocTensor("features_left_08", feat_08_dims_);
|
||||
d_feat_left_16_ = allocTensor("features_left_16", feat_16_dims_);
|
||||
d_feat_left_32_ = allocTensor("features_left_32", feat_32_dims_);
|
||||
d_feat_right_04_ = allocTensor("features_right_04", feat_04_dims_);
|
||||
d_stem_2x_ = allocTensor("stem_2x", stem_2x_dims_);
|
||||
|
||||
feature_context_->setTensorAddress("left", d_left_);
|
||||
feature_context_->setTensorAddress("right", d_right_);
|
||||
feature_context_->setTensorAddress("features_left_04", d_feat_left_04_);
|
||||
feature_context_->setTensorAddress("features_left_08", d_feat_left_08_);
|
||||
feature_context_->setTensorAddress("features_left_16", d_feat_left_16_);
|
||||
feature_context_->setTensorAddress("features_left_32", d_feat_left_32_);
|
||||
feature_context_->setTensorAddress("features_right_04", d_feat_right_04_);
|
||||
feature_context_->setTensorAddress("stem_2x", d_stem_2x_);
|
||||
}
|
||||
|
||||
void FFSDepthInference::allocatePostBuffers() {
|
||||
const size_t H = static_cast<size_t>(config_.image_height);
|
||||
const size_t W = static_cast<size_t>(config_.image_width);
|
||||
const size_t H4 = H / 4;
|
||||
const size_t W4 = W / 4;
|
||||
|
||||
const auto gwc_dt = post_engine_->getTensorDataType("gwc_volume");
|
||||
gwc_fp16_ = (gwc_dt == nvinfer1::DataType::kHALF);
|
||||
const size_t gwc_elem = gwc_fp16_ ? sizeof(__half) : sizeof(float);
|
||||
const size_t gwc_bytes =
|
||||
static_cast<size_t>(config_.cv_group) *
|
||||
static_cast<size_t>(gwc_disp_levels_) * H4 * W4 * gwc_elem;
|
||||
cudaMallocChecked(reinterpret_cast<void**>(&d_gwc_volume_), gwc_bytes, "gwc_volume");
|
||||
|
||||
const auto disp_dt = post_engine_->getTensorDataType("disp");
|
||||
const size_t disp_bytes = H * W * elementSize(disp_dt);
|
||||
cudaMallocChecked(reinterpret_cast<void**>(&d_disp_), disp_bytes, "d_disp_");
|
||||
|
||||
const size_t disp_cropped_bytes = H * W * sizeof(float);
|
||||
cudaMallocChecked(reinterpret_cast<void**>(&d_disp_cropped_),
|
||||
disp_cropped_bytes, "d_disp_cropped_");
|
||||
|
||||
post_context_->setTensorAddress("features_left_04", d_feat_left_04_);
|
||||
post_context_->setTensorAddress("features_left_08", d_feat_left_08_);
|
||||
post_context_->setTensorAddress("features_left_16", d_feat_left_16_);
|
||||
post_context_->setTensorAddress("features_left_32", d_feat_left_32_);
|
||||
post_context_->setTensorAddress("features_right_04", d_feat_right_04_);
|
||||
post_context_->setTensorAddress("stem_2x", d_stem_2x_);
|
||||
post_context_->setTensorAddress("gwc_volume", d_gwc_volume_);
|
||||
post_context_->setTensorAddress("disp", d_disp_);
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Pre-processing
|
||||
// =========================================================================
|
||||
|
||||
void FFSDepthInference::preprocessRGBGPU(
|
||||
const uint8_t* d_rgb, int src_h, int src_w, float* d_output)
|
||||
{
|
||||
const int mH = config_.image_height;
|
||||
const int mW = config_.image_width;
|
||||
|
||||
if (src_h == mH && src_w == mW) {
|
||||
scaled_h_ = mH;
|
||||
scaled_w_ = mW;
|
||||
cuda::ffsCudaPreprocessRGBToCHW(
|
||||
d_rgb, d_output, src_h, src_w, mH, mW, stream_);
|
||||
} else {
|
||||
const float scale = std::min(static_cast<float>(mW) / src_w,
|
||||
static_cast<float>(mH) / src_h);
|
||||
scaled_w_ = std::max(1, static_cast<int>(std::round(src_w * scale)));
|
||||
scaled_h_ = std::max(1, static_cast<int>(std::round(src_h * scale)));
|
||||
cuda::ffsCudaResizeUniformAndPad(
|
||||
d_rgb, d_output, src_h, src_w, scaled_h_, scaled_w_, mH, mW, stream_);
|
||||
}
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Pipeline stages
|
||||
// =========================================================================
|
||||
|
||||
void FFSDepthInference::runFeatureRunner() {
|
||||
if (!feature_context_->enqueueV3(stream_))
|
||||
throw std::runtime_error("[FFS] feature_runner failed");
|
||||
}
|
||||
|
||||
void FFSDepthInference::buildGWCVolume() {
|
||||
int B = feat_04_dims_[0];
|
||||
int C = feat_04_dims_[1];
|
||||
int H = feat_04_dims_[2];
|
||||
int W = feat_04_dims_[3];
|
||||
|
||||
// Only mixed precision path is used (FP32 features -> FP16 GWC volume)
|
||||
cuda::ffsCudaBuildGWCVolumeMixed(
|
||||
d_feat_left_04_, d_feat_right_04_,
|
||||
reinterpret_cast<__half*>(d_gwc_volume_),
|
||||
B, C, H, W, gwc_disp_levels_, config_.cv_group, config_.normalize, stream_);
|
||||
|
||||
}
|
||||
|
||||
void FFSDepthInference::runPostRunner() {
|
||||
if (!post_context_->enqueueV3(stream_))
|
||||
throw std::runtime_error("[FFS] post_runner failed");
|
||||
}
|
||||
|
||||
|
||||
// =========================================================================
|
||||
// Public inference entry point
|
||||
// =========================================================================
|
||||
|
||||
void FFSDepthInference::infer(
|
||||
const uint8_t* d_left_rgb, const uint8_t* d_right_rgb,
|
||||
int input_h, int input_w,
|
||||
float* d_disp_out)
|
||||
{
|
||||
if (!d_left_rgb || !d_right_rgb || !d_disp_out) {
|
||||
throw std::runtime_error("[FFS] infer: null device pointer");
|
||||
}
|
||||
if (input_h <= 0 || input_w <= 0) {
|
||||
throw std::runtime_error("[FFS] infer: input dimensions must be positive");
|
||||
}
|
||||
const int mH = config_.image_height;
|
||||
const int mW = config_.image_width;
|
||||
const bool needs_resize = (input_h != mH || input_w != mW);
|
||||
preprocessRGBGPU(d_left_rgb, input_h, input_w, d_left_);
|
||||
preprocessRGBGPU(d_right_rgb, input_h, input_w, d_right_);
|
||||
runFeatureRunner();
|
||||
buildGWCVolume();
|
||||
if (!needs_resize) {
|
||||
post_context_->setTensorAddress("disp", d_disp_out);
|
||||
}
|
||||
runPostRunner();
|
||||
if (needs_resize) {
|
||||
cuda::ffsCudaClampDisparity(d_disp_, mH * mW, 0.0f, stream_);
|
||||
// Crop padding: model (mW x mH) -> scaled (scaled_w_ x scaled_h_)
|
||||
cuda::ffsCudaCropDisparity(
|
||||
d_disp_, d_disp_cropped_,
|
||||
mH, mW, scaled_h_, scaled_w_, stream_);
|
||||
// Upsample to input resolution with disparity scale correction
|
||||
float disp_scale = static_cast<float>(input_w) / scaled_w_;
|
||||
cuda::ffsCudaUpsampleDisparity(
|
||||
d_disp_cropped_, d_disp_out,
|
||||
scaled_w_, scaled_h_, input_w, input_h,
|
||||
disp_scale, stream_);
|
||||
} else {
|
||||
cuda::ffsCudaClampDisparity(d_disp_out, mH * mW, 0.0f, stream_);
|
||||
post_context_->setTensorAddress("disp", d_disp_);
|
||||
}
|
||||
}
|
||||
|
||||
void FFSDepthInference::sync() {
|
||||
cudaStreamSynchronize(stream_);
|
||||
}
|
||||
|
||||
void FFSDepthInference::dispToDepth(
|
||||
const float* d_disp,
|
||||
int height, int width,
|
||||
float fx, float baseline_m,
|
||||
float* d_depth_out)
|
||||
{
|
||||
cuda::ffsCudaDispToDepth(d_disp, d_depth_out,
|
||||
height, width, fx, baseline_m, stream_);
|
||||
}
|
||||
|
||||
void FFSDepthInference::inferDepth(
|
||||
const uint8_t* d_left_rgb, const uint8_t* d_right_rgb,
|
||||
int input_h, int input_w,
|
||||
float fx, float baseline_m,
|
||||
float* d_depth_out)
|
||||
{
|
||||
if (input_h <= 0 || input_w <= 0) {
|
||||
throw std::runtime_error("[FFS] inferDepth: input dimensions must be positive");
|
||||
}
|
||||
const size_t num_pixels = static_cast<size_t>(input_h) * static_cast<size_t>(input_w);
|
||||
|
||||
if (static_cast<int64_t>(num_pixels) > depth_alloc_pixels_) {
|
||||
// cudaFree is implicitly stream-ordered: it waits for prior work on the
|
||||
// device before reclaiming the allocation, so this is safe to call here
|
||||
// even though earlier inferDepth() calls may still be in flight.
|
||||
if (d_disp_for_depth_) {
|
||||
cudaFree(d_disp_for_depth_);
|
||||
d_disp_for_depth_ = nullptr;
|
||||
}
|
||||
cudaMallocChecked(reinterpret_cast<void**>(&d_disp_for_depth_),
|
||||
num_pixels * sizeof(float), "d_disp_for_depth_");
|
||||
depth_alloc_pixels_ = static_cast<int64_t>(num_pixels);
|
||||
}
|
||||
|
||||
infer(d_left_rgb, d_right_rgb, input_h, input_w, d_disp_for_depth_);
|
||||
cuda::ffsCudaDispToDepth(d_disp_for_depth_, d_depth_out,
|
||||
input_h, input_w, fx, baseline_m, stream_);
|
||||
}
|
||||
|
||||
} // namespace ffs_depth
|
||||
@@ -0,0 +1,325 @@
|
||||
#include "ffs_gwc_plugin.hpp"
|
||||
|
||||
#include <NvInfer.h>
|
||||
#include <NvInferRuntimePlugin.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <cstring>
|
||||
#include <iostream>
|
||||
#include <mutex>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace ffs_depth {
|
||||
namespace cuda {
|
||||
extern "C" {
|
||||
|
||||
void ffsCudaBuildGWCVolumeFloat(
|
||||
const float* d_fl, const float* d_fr, float* d_gwc,
|
||||
int B, int C, int H, int W, int max_disp, int ngroups, bool normalize, cudaStream_t s);
|
||||
|
||||
void ffsCudaBuildGWCVolumeHalf(
|
||||
const __half* d_fl, const __half* d_fr, __half* d_gwc,
|
||||
int B, int C, int H, int W, int max_disp, int ngroups, bool normalize, cudaStream_t s);
|
||||
|
||||
void ffsCudaBuildGWCVolumeHalfToFloat(
|
||||
const __half* d_fl, const __half* d_fr, float* d_gwc,
|
||||
int B, int C, int H, int W, int max_disp, int ngroups, bool normalize, cudaStream_t s);
|
||||
|
||||
void ffsCudaBuildGWCVolumeMixed(
|
||||
const float* d_fl, const float* d_fr, __half* d_gwc,
|
||||
int B, int C, int H, int W, int max_disp, int ngroups, bool normalize, cudaStream_t s);
|
||||
|
||||
} // extern "C"
|
||||
} // namespace cuda
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr char kPluginName[] = "FFSGWCVolume";
|
||||
constexpr char kPluginVersion[] = "1";
|
||||
|
||||
struct GWCParams {
|
||||
int32_t max_disp = 0; // Disparity levels at feature resolution, e.g. max_disp / 4.
|
||||
int32_t cv_group = 0;
|
||||
int32_t normalize = 1;
|
||||
};
|
||||
|
||||
int32_t fieldToInt(nvinfer1::PluginField const& field, int32_t fallback) {
|
||||
if (!field.data || field.length <= 0) return fallback;
|
||||
if (field.type == nvinfer1::PluginFieldType::kINT32) {
|
||||
return *static_cast<int32_t const*>(field.data);
|
||||
}
|
||||
if (field.type == nvinfer1::PluginFieldType::kINT64) {
|
||||
return static_cast<int32_t>(*static_cast<int64_t const*>(field.data));
|
||||
}
|
||||
return fallback;
|
||||
}
|
||||
|
||||
class FFSGWCVolumePlugin final : public nvinfer1::IPluginV2DynamicExt {
|
||||
public:
|
||||
explicit FFSGWCVolumePlugin(GWCParams params) : params_(params) {}
|
||||
|
||||
FFSGWCVolumePlugin(void const* data, size_t length) {
|
||||
if (data && length == sizeof(GWCParams)) {
|
||||
std::memcpy(¶ms_, data, sizeof(GWCParams));
|
||||
}
|
||||
}
|
||||
|
||||
nvinfer1::IPluginV2DynamicExt* clone() const noexcept override {
|
||||
auto* plugin = new FFSGWCVolumePlugin(params_);
|
||||
plugin->setPluginNamespace(namespace_.c_str());
|
||||
return plugin;
|
||||
}
|
||||
|
||||
char const* getPluginType() const noexcept override { return kPluginName; }
|
||||
char const* getPluginVersion() const noexcept override { return kPluginVersion; }
|
||||
int32_t getNbOutputs() const noexcept override { return 1; }
|
||||
|
||||
nvinfer1::DimsExprs getOutputDimensions(
|
||||
int32_t outputIndex,
|
||||
nvinfer1::DimsExprs const* inputs,
|
||||
int32_t nbInputs,
|
||||
nvinfer1::IExprBuilder& exprBuilder) noexcept override {
|
||||
nvinfer1::DimsExprs out{};
|
||||
if (outputIndex != 0 || nbInputs != 2 || inputs[0].nbDims != 4) {
|
||||
out.nbDims = -1;
|
||||
return out;
|
||||
}
|
||||
|
||||
out.nbDims = 5;
|
||||
out.d[0] = inputs[0].d[0]; // B
|
||||
out.d[1] = exprBuilder.constant(params_.cv_group);
|
||||
out.d[2] = exprBuilder.constant(params_.max_disp);
|
||||
out.d[3] = inputs[0].d[2]; // H
|
||||
out.d[4] = inputs[0].d[3]; // W
|
||||
return out;
|
||||
}
|
||||
|
||||
bool supportsFormatCombination(
|
||||
int32_t pos,
|
||||
nvinfer1::PluginTensorDesc const* inOut,
|
||||
int32_t nbInputs,
|
||||
int32_t nbOutputs) noexcept override {
|
||||
if (nbInputs != 2 || nbOutputs != 1 || pos < 0 || pos >= 3) return false;
|
||||
if (inOut[pos].format != nvinfer1::TensorFormat::kLINEAR) return false;
|
||||
|
||||
if (pos == 0) {
|
||||
return inOut[0].type == nvinfer1::DataType::kFLOAT ||
|
||||
inOut[0].type == nvinfer1::DataType::kHALF;
|
||||
}
|
||||
if (pos == 1) {
|
||||
return inOut[1].type == inOut[0].type;
|
||||
}
|
||||
return inOut[2].type == nvinfer1::DataType::kFLOAT ||
|
||||
inOut[2].type == nvinfer1::DataType::kHALF;
|
||||
}
|
||||
|
||||
void configurePlugin(
|
||||
nvinfer1::DynamicPluginTensorDesc const*,
|
||||
int32_t,
|
||||
nvinfer1::DynamicPluginTensorDesc const*,
|
||||
int32_t) noexcept override {}
|
||||
|
||||
size_t getWorkspaceSize(
|
||||
nvinfer1::PluginTensorDesc const*,
|
||||
int32_t,
|
||||
nvinfer1::PluginTensorDesc const*,
|
||||
int32_t) const noexcept override {
|
||||
return 0;
|
||||
}
|
||||
|
||||
int32_t enqueue(
|
||||
nvinfer1::PluginTensorDesc const* inputDesc,
|
||||
nvinfer1::PluginTensorDesc const* outputDesc,
|
||||
void const* const* inputs,
|
||||
void* const* outputs,
|
||||
void*,
|
||||
cudaStream_t stream) noexcept override {
|
||||
if (!inputs || !outputs || !inputs[0] || !inputs[1] || !outputs[0]) return 1;
|
||||
if (inputDesc[0].dims.nbDims != 4) return 1;
|
||||
|
||||
const int B = inputDesc[0].dims.d[0];
|
||||
const int C = inputDesc[0].dims.d[1];
|
||||
const int H = inputDesc[0].dims.d[2];
|
||||
const int W = inputDesc[0].dims.d[3];
|
||||
const bool normalize = params_.normalize != 0;
|
||||
|
||||
if (inputDesc[0].type == nvinfer1::DataType::kFLOAT &&
|
||||
outputDesc[0].type == nvinfer1::DataType::kFLOAT) {
|
||||
cuda::ffsCudaBuildGWCVolumeFloat(
|
||||
static_cast<float const*>(inputs[0]),
|
||||
static_cast<float const*>(inputs[1]),
|
||||
static_cast<float*>(outputs[0]),
|
||||
B, C, H, W, params_.max_disp, params_.cv_group, normalize, stream);
|
||||
return 0;
|
||||
}
|
||||
if (inputDesc[0].type == nvinfer1::DataType::kFLOAT &&
|
||||
outputDesc[0].type == nvinfer1::DataType::kHALF) {
|
||||
cuda::ffsCudaBuildGWCVolumeMixed(
|
||||
static_cast<float const*>(inputs[0]),
|
||||
static_cast<float const*>(inputs[1]),
|
||||
static_cast<__half*>(outputs[0]),
|
||||
B, C, H, W, params_.max_disp, params_.cv_group, normalize, stream);
|
||||
return 0;
|
||||
}
|
||||
if (inputDesc[0].type == nvinfer1::DataType::kHALF &&
|
||||
outputDesc[0].type == nvinfer1::DataType::kHALF) {
|
||||
cuda::ffsCudaBuildGWCVolumeHalf(
|
||||
static_cast<__half const*>(inputs[0]),
|
||||
static_cast<__half const*>(inputs[1]),
|
||||
static_cast<__half*>(outputs[0]),
|
||||
B, C, H, W, params_.max_disp, params_.cv_group, normalize, stream);
|
||||
return 0;
|
||||
}
|
||||
if (inputDesc[0].type == nvinfer1::DataType::kHALF &&
|
||||
outputDesc[0].type == nvinfer1::DataType::kFLOAT) {
|
||||
cuda::ffsCudaBuildGWCVolumeHalfToFloat(
|
||||
static_cast<__half const*>(inputs[0]),
|
||||
static_cast<__half const*>(inputs[1]),
|
||||
static_cast<float*>(outputs[0]),
|
||||
B, C, H, W, params_.max_disp, params_.cv_group, normalize, stream);
|
||||
return 0;
|
||||
}
|
||||
|
||||
return 1;
|
||||
}
|
||||
|
||||
nvinfer1::DataType getOutputDataType(
|
||||
int32_t,
|
||||
nvinfer1::DataType const* inputTypes,
|
||||
int32_t nbInputs) const noexcept override {
|
||||
if (nbInputs > 0 && inputTypes[0] == nvinfer1::DataType::kHALF) {
|
||||
return nvinfer1::DataType::kHALF;
|
||||
}
|
||||
return nvinfer1::DataType::kFLOAT;
|
||||
}
|
||||
|
||||
int32_t initialize() noexcept override { return 0; }
|
||||
void terminate() noexcept override {}
|
||||
size_t getSerializationSize() const noexcept override { return sizeof(GWCParams); }
|
||||
void serialize(void* buffer) const noexcept override {
|
||||
std::memcpy(buffer, ¶ms_, sizeof(GWCParams));
|
||||
}
|
||||
void destroy() noexcept override { delete this; }
|
||||
void setPluginNamespace(char const* pluginNamespace) noexcept override {
|
||||
namespace_ = pluginNamespace ? pluginNamespace : "";
|
||||
}
|
||||
char const* getPluginNamespace() const noexcept override { return namespace_.c_str(); }
|
||||
|
||||
void attachToContext(cudnnContext*, cublasContext*, nvinfer1::IGpuAllocator*) noexcept override {}
|
||||
void detachFromContext() noexcept override {}
|
||||
|
||||
private:
|
||||
GWCParams params_;
|
||||
std::string namespace_;
|
||||
};
|
||||
|
||||
class FFSGWCVolumePluginCreator final : public nvinfer1::IPluginCreator {
|
||||
public:
|
||||
FFSGWCVolumePluginCreator() {
|
||||
fields_.emplace_back("max_disp", nullptr, nvinfer1::PluginFieldType::kINT32, 1);
|
||||
fields_.emplace_back("cv_group", nullptr, nvinfer1::PluginFieldType::kINT32, 1);
|
||||
fields_.emplace_back("normalize", nullptr, nvinfer1::PluginFieldType::kINT32, 1);
|
||||
field_collection_.nbFields = static_cast<int32_t>(fields_.size());
|
||||
field_collection_.fields = fields_.data();
|
||||
}
|
||||
|
||||
char const* getPluginName() const noexcept override { return kPluginName; }
|
||||
char const* getPluginVersion() const noexcept override { return kPluginVersion; }
|
||||
nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override {
|
||||
return &field_collection_;
|
||||
}
|
||||
|
||||
nvinfer1::IPluginV2* createPlugin(
|
||||
char const*,
|
||||
nvinfer1::PluginFieldCollection const* fc) noexcept override {
|
||||
GWCParams params;
|
||||
if (fc) {
|
||||
for (int32_t i = 0; i < fc->nbFields; ++i) {
|
||||
auto const& field = fc->fields[i];
|
||||
if (!std::strcmp(field.name, "max_disp")) {
|
||||
params.max_disp = fieldToInt(field, params.max_disp);
|
||||
} else if (!std::strcmp(field.name, "cv_group")) {
|
||||
params.cv_group = fieldToInt(field, params.cv_group);
|
||||
} else if (!std::strcmp(field.name, "normalize")) {
|
||||
params.normalize = fieldToInt(field, params.normalize);
|
||||
}
|
||||
}
|
||||
}
|
||||
if (params.max_disp <= 0 || params.cv_group <= 0) return nullptr;
|
||||
auto* plugin = new FFSGWCVolumePlugin(params);
|
||||
plugin->setPluginNamespace(namespace_.c_str());
|
||||
return plugin;
|
||||
}
|
||||
|
||||
nvinfer1::IPluginV2* deserializePlugin(
|
||||
char const*,
|
||||
void const* serialData,
|
||||
size_t serialLength) noexcept override {
|
||||
auto* plugin = new FFSGWCVolumePlugin(serialData, serialLength);
|
||||
plugin->setPluginNamespace(namespace_.c_str());
|
||||
return plugin;
|
||||
}
|
||||
|
||||
void setPluginNamespace(char const* pluginNamespace) noexcept override {
|
||||
namespace_ = pluginNamespace ? pluginNamespace : "";
|
||||
}
|
||||
char const* getPluginNamespace() const noexcept override {
|
||||
return namespace_.c_str();
|
||||
}
|
||||
|
||||
private:
|
||||
std::vector<nvinfer1::PluginField> fields_;
|
||||
nvinfer1::PluginFieldCollection field_collection_{};
|
||||
std::string namespace_;
|
||||
};
|
||||
|
||||
FFSGWCVolumePluginCreator g_creator;
|
||||
|
||||
} // namespace
|
||||
|
||||
bool registerFFSGWCPlugin() {
|
||||
static bool result = false;
|
||||
static std::once_flag once;
|
||||
std::call_once(once, []() {
|
||||
auto try_register = [](nvinfer1::IPluginRegistry* registry, const char* tag) -> bool {
|
||||
if (!registry) {
|
||||
std::cerr << "[FFS plugin] " << tag << ": registry unavailable\n";
|
||||
return false;
|
||||
}
|
||||
if (registry->getPluginCreator(kPluginName, kPluginVersion, "")) {
|
||||
std::cerr << "[FFS plugin] " << tag
|
||||
<< ": FFSGWCVolume v1 ns=\"\" found-already\n";
|
||||
return true;
|
||||
}
|
||||
if (registry->registerCreator(g_creator, "")) {
|
||||
std::cerr << "[FFS plugin] " << tag
|
||||
<< ": FFSGWCVolume v1 ns=\"\" newly-registered\n";
|
||||
return true;
|
||||
}
|
||||
std::cerr << "[FFS plugin] " << tag
|
||||
<< ": FFSGWCVolume v1 ns=\"\" register FAILED\n";
|
||||
return false;
|
||||
};
|
||||
|
||||
const bool a = try_register(::getPluginRegistry(), "runtime-registry");
|
||||
const bool b = try_register(
|
||||
nvinfer1::getBuilderPluginRegistry(nvinfer1::EngineCapability::kSTANDARD),
|
||||
"builder-registry");
|
||||
result = a || b;
|
||||
});
|
||||
return result;
|
||||
}
|
||||
|
||||
// Auto-register the creator with the global TRT plugin registry on shared
|
||||
// library load. Lets external tools (trtexec --staticPlugins, polygraphy
|
||||
// --plugins) deserialize engines containing FFSGWCVolume without calling
|
||||
// registerFFSGWCPlugin() themselves.
|
||||
REGISTER_TENSORRT_PLUGIN(FFSGWCVolumePluginCreator);
|
||||
|
||||
} // namespace ffs_depth
|
||||
|
||||
extern "C" bool ffs_register_gwc_plugin() {
|
||||
return ffs_depth::registerFFSGWCPlugin();
|
||||
}
|
||||
Reference in New Issue
Block a user