Răsfoiți Sursa

add depth result

leon 1 lună în urmă
părinte
comite
936e568565

+ 2 - 1
src/infer/trt/depth/depth.cu

@@ -140,7 +140,8 @@ Result DepthModelImpl::forward(const tensor::Image &image, void *stream)
     checkRuntime(cudaMemcpyAsync(depth_map_buffer_.cpu(), depth_map_buffer_.gpu(),
                                     depth_map_buffer_.gpu_bytes(), cudaMemcpyDeviceToHost, stream_));
     checkRuntime(cudaStreamSynchronize(stream_));
-    return cv::Mat();
+    cv::Mat depth_mat(image.height, image.width, CV_32FC1, depth_map_buffer_.cpu());
+    return depth_mat;
 }
 
 

+ 1 - 0
src/nodes/draw/drawNode.cpp

@@ -52,6 +52,7 @@ void DrawNode::work()
                 // printf("Node %s push data to %s\n", name_.c_str(), output_buffer.first.c_str());
                 output_buffer.second->push(metaData);
             }
+            printf(metaData->depth.size() > 0 ? "depth size: %d\n" : "depth size: 0\n", metaData->depth.size());
         }
         if (!has_data) {
             printf("draw wait data");

+ 7 - 3
src/nodes/infer/inferNode.cpp

@@ -33,7 +33,8 @@ void InferNode::work()
             }
             auto det_result = model_->forward(tensor::cvimg(image));
 
-            if (std::holds_alternative<data::BoxArray>(det_result)) {
+            if (std::holds_alternative<data::BoxArray>(det_result)) 
+            {
                 auto result = std::get<data::BoxArray>(det_result);
                 for (auto& r : result)
                 {
@@ -41,9 +42,12 @@ void InferNode::work()
                 }
                 metaData->boxes = result;
                 // 处理检测框...
-            } else if(std::holds_alternative<cv::Mat>(det_result)){
+            } 
+            else if(std::holds_alternative<cv::Mat>(det_result))
+            {
                 // 记录错误或抛出明确异常
-                auto result = std::get<cv::Mat>(det_result);
+                auto depth_mat = std::get<cv::Mat>(det_result);
+                metaData->depth = depth_mat;
             }
             else
             {