leon 1 miesiąc temu
rodzic
commit
b3925253ef
2 zmienionych plików z 16 dodań i 4 usunięć
  1. 6 4
      src/infer/trt/depth/depth.cu
  2. 10 0
      src/main.cpp

+ 6 - 4
src/infer/trt/depth/depth.cu

@@ -118,14 +118,14 @@ Result DepthModelImpl::forward(const tensor::Image &image, void *stream)
     #ifdef TRT10
     if (!trt_->forward(std::unordered_map<std::string, const void *>{
             { "input", input_buffer_.gpu() }, 
-            { "output", depth_map_buffer_.gpu() }
+            { "output", output_buffer_.gpu() }
         }, stream_))
     {
-        printf("Failed to tensorRT forward.");
-        return {};
+        printf("Failed to tensorRT forward.\n");
+        return cv::Mat();
     }
     #else
-    std::vector<void *> bindings{input_buffer_.gpu(), depth_map_buffer_.gpu()};
+    std::vector<void *> bindings{input_buffer_.gpu(), output_buffer_.gpu()};
     if (!trt_->forward(bindings, stream)) 
     {
         printf("Failed to tensorRT forward.");
@@ -133,6 +133,8 @@ Result DepthModelImpl::forward(const tensor::Image &image, void *stream)
     }
     #endif
 
+    adjust_memory(image.width, image.height);
+
     postprocess(image.width, image.height, stream);
 
     checkRuntime(cudaMemcpyAsync(depth_map_buffer_.cpu(), depth_map_buffer_.gpu(),

+ 10 - 0
src/main.cpp

@@ -27,6 +27,16 @@ void test_depth()
     getchar();
 }
 
+void test_yolo()
+{
+
+}
+
+void test_multi_model()
+{
+    
+}
+
 int main()
 {
     test_depth();