From 5a2903363507ab3fc720c45c282b153806657917 Mon Sep 17 00:00:00 2001
From: luoyao <luoyao@baidu.com>
Date: Tue, 19 Dec 2023 17:10:18 +0800
Subject: [PATCH] update

---
 .../metric3d_benchmark.cpp                    | 31 +++++++++++++------
 src/models/model_io_define.h                  |  1 +
 src/models/mono_depth_estimation/metric3d.inl |  4 +++
 3 files changed, 26 insertions(+), 10 deletions(-)

diff --git a/src/apps/model_benchmark/mono_depth_estimation/metric3d_benchmark.cpp b/src/apps/model_benchmark/mono_depth_estimation/metric3d_benchmark.cpp
index a714116a..e78c69d6 100644
--- a/src/apps/model_benchmark/mono_depth_estimation/metric3d_benchmark.cpp
+++ b/src/apps/model_benchmark/mono_depth_estimation/metric3d_benchmark.cpp
@@ -79,17 +79,28 @@ int main(int argc, char** argv) {
     LOG(INFO) << "cost time: " << cost_time << "s, fps: " << loop_times / cost_time;
 
     // save colorized depth map
-    cv::Mat colorized_depth_map;
-    CvUtils::colorize_depth_map(model_output.depth_map, colorized_depth_map);
     std::string output_file_name = FilePathUtil::get_file_name(input_image_path);
-    output_file_name = output_file_name.substr(0, output_file_name.find_last_of('.')) + "_metric3d_depth_result.png";
-    std::string output_path = FilePathUtil::concat_path(
-        "../demo_data/model_test_input/mono_depth_estimation", output_file_name);
-    cv::imwrite(output_path, colorized_depth_map);
-    LOG(INFO) << "prediction result image has been written into: " << output_path;
-
-    return 0;
-}
+    output_file_name = output_file_name.substr(0, output_file_name.find_last_of('.')) + "_metric3d_colorized_depth_result.png";
+    std::string output_path = FilePathUtil::concat_path("../demo_data/model_test_input/mono_depth_estimation", output_file_name);
+    cv::imwrite(output_path, model_output.colorized_depth_map);
+    LOG(INFO) << "prediction colorized depth image has been written into: " << output_path;
 
+    // save depth map and confidence map
+    output_file_name = FilePathUtil::get_file_name(input_image_path);
+    output_file_name = output_file_name.substr(0, output_file_name.find_last_of('.')) + "_metric3d_depth_result.yaml";
+    output_path = FilePathUtil::concat_path("../demo_data/model_test_input/mono_depth_estimation", output_file_name);
+    cv::FileStorage out_depth_map;
+    out_depth_map.open(output_path, cv::FileStorage::WRITE);
+    out_depth_map.write("depth_map", model_output.depth_map);
+    LOG(INFO) << "prediction depth map has been written into: " << output_path;
 
+    output_file_name = FilePathUtil::get_file_name(input_image_path);
+    output_file_name = output_file_name.substr(0, output_file_name.find_last_of('.')) + "_metric3d_conf_result.yaml";
+    output_path = FilePathUtil::concat_path("../demo_data/model_test_input/mono_depth_estimation", output_file_name);
+    cv::FileStorage out_conf_map;
+    out_conf_map.open(output_path, cv::FileStorage::WRITE);
+    out_conf_map.write("conf_map", model_output.confidence_map);
+    LOG(INFO) << "prediction confidence map has been written into: " << output_path;
 
+    return 0;
+}
diff --git a/src/models/model_io_define.h b/src/models/model_io_define.h
index 9e9d4a12..f2b2a745 100644
--- a/src/models/model_io_define.h
+++ b/src/models/model_io_define.h
@@ -133,6 +133,7 @@ using std_feature_point_match_output = std::vector<matched_fp>;
 namespace mono_depth_estimation {
 
 struct mde_output {
+    cv::Mat confidence_map;
     cv::Mat depth_map;
     cv::Mat colorized_depth_map;
 };
diff --git a/src/models/mono_depth_estimation/metric3d.inl b/src/models/mono_depth_estimation/metric3d.inl
index 628a1ad9..8c8e21b8 100644
--- a/src/models/mono_depth_estimation/metric3d.inl
+++ b/src/models/mono_depth_estimation/metric3d.inl
@@ -113,7 +113,9 @@ template <typename OUTPUT>
 typename std::enable_if<std::is_same<OUTPUT, std::decay<std_mde_output>::type>::value, std_mde_output>::type
 transform_output(const metric3d_impl::internal_output& internal_out) {
     std_mde_output result;
+    internal_out.confidence_map.copyTo(result.confidence_map);
     internal_out.depth_map.copyTo(result.depth_map);
+    internal_out.colorized_depth_map.copyTo(result.colorized_depth_map);
     return result;
 }
 
@@ -653,6 +655,7 @@ metric3d_impl::internal_output Metric3D<INPUT, OUTPUT>::Impl::mnn_decode_output(
 
     // copy result
     std_mde_output out;
+    confidence_map.copyTo(out.confidence_map);
     depth_map.copyTo(out.depth_map);
     colorized_depth_map.copyTo(out.colorized_depth_map);
     return out;
@@ -906,6 +909,7 @@ metric3d_impl::internal_output Metric3D<INPUT, OUTPUT>::Impl::trt_decode_output(
 
     // copy result
     std_mde_output out;
+    confidence_map.copyTo(out.confidence_map);
     depth_map.copyTo(out.depth_map);
     colorized_depth_map.copyTo(out.colorized_depth_map);
     return out;