leon преди 5 месеца
родител
ревизия
78a8ffc11e
променени са 2 файла, в които са добавени 9 реда и са изтрити 2 реда
  1. 7 0
      src/interface.cpp
  2. 2 2
      test.py

+ 7 - 0
src/interface.cpp

@@ -62,6 +62,12 @@ public:
         return instance_->forward(cvimg(image));
     }
 
+    resnet::Attribute forward_path(const std::string& image_path)
+    {
+        cv::Mat image = cv::imread(image_path);
+        return instance_->forward(cvimg(image));
+    }
+
 
     bool valid(){
 		return instance_ != nullptr;
@@ -86,6 +92,7 @@ PYBIND11_MODULE(trtresnet, m){
 	py::class_<TrtResnetInfer>(m, "TrtResnetInfer")
 		.def(py::init<string>(), py::arg("model_path"))
 		.def_property_readonly("valid", &TrtResnetInfer::valid)
+        .def("forward_path", &TrtResnetInfer::forward_path, py::arg("image_path"))
 		.def("forward", &TrtResnetInfer::forward, py::arg("image"));
 };
 

+ 2 - 2
test.py

@@ -3,7 +3,7 @@ import cv2
 
 infer = trtresnet.TrtResnetInfer("workspace/resnet.engine")
 
-image = cv2.imread("workspace/inference/car.jpg")
-result = infer.forward(image)
+# image = cv2.imread("workspace/inference/car.jpg")
+result = infer.forward_path("workspace/inference/car.jpg")
 
 print(result)