Добавлены пропсы конвейера и стереодвижки, задействованные в прогоне
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,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;
|
||||
}
|
||||
Reference in New Issue
Block a user