add opencv gpu support

Signed-off-by: jagadeesh <jagadeeshj@ideas2it.com>
pytorch · Aug 30, 2023 · d4a431a · d4a431a
1 parent 46fadcd
commit d4a431a
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 9 deletions.
diff --git a/cpp/src/examples/CMakeLists.txt b/cpp/src/examples/CMakeLists.txt
@@ -17,3 +17,6 @@ target_include_directories(resnet-18_handler PUBLIC ${OPENCV_DIR})
 target_include_directories(resnet-18_handler PUBLIC ${RESNET_SRC_DIR})
 target_link_libraries(resnet-18_handler PRIVATE ts_backends_torch_scripted ts_utils ${TORCH_LIBRARIES})
 target_link_libraries(resnet-18_handler PRIVATE "/usr/local/lib/libopencv_imgcodecs.so")
+target_link_libraries(resnet-18_handler PRIVATE "/usr/local/lib/libopencv_cudawarping.so")
+target_link_libraries(resnet-18_handler PRIVATE "/usr/local/lib/libopencv_cudaimgproc.so")
+target_link_libraries(resnet-18_handler PRIVATE "/usr/local/lib/libopencv_core.so")
diff --git a/cpp/src/examples/image_classifier/resnet-18/resnet-18_handler.cc b/cpp/src/examples/image_classifier/resnet-18/resnet-18_handler.cc
@@ -3,6 +3,8 @@
 #include <folly/json.h>
 
 #include <fstream>
+#include <opencv2/cudaimgproc.hpp>
+#include <opencv2/cudawarping.hpp>
 #include <opencv2/opencv.hpp>
 
 namespace resnet {
@@ -88,7 +90,6 @@ std::vector<torch::jit::IValue> ResnetHandler::Preprocess(
           std::cerr << "Failed to decode the image.\n";
         }
 
-        // Crop image
         const int rows = image.rows;
         const int cols = image.cols;
 
@@ -99,27 +100,39 @@ std::vector<torch::jit::IValue> ResnetHandler::Preprocess(
         const cv::Rect roi(offsetW, offsetH, cropSize, cropSize);
         image = image(roi);
 
-        // Resize
-        cv::resize(image, image, cv::Size(kTargetImageSize, kTargetImageSize));
+        // Convert the image to GPU Mat
+        cv::cuda::GpuMat gpuImage;
+        cv::Mat resultImage;
 
-        // Convert BGR to RGB format
-        cv::cvtColor(image, image, cv::COLOR_BGR2RGB);
+        gpuImage.upload(image);
 
-        image.convertTo(image, CV_32FC3, 1 / 255.0);
+        // Resize on GPU
+        cv::cuda::resize(gpuImage, gpuImage,
+                         cv::Size(kTargetImageSize, kTargetImageSize));
 
-        // Convert the OpenCV image to a torch tensor
+        // Convert to BGR on GPU
+        cv::cuda::cvtColor(gpuImage, gpuImage, cv::COLOR_BGR2RGB);
+
+        // Convert to float on GPU
+        gpuImage.convertTo(gpuImage, CV_32FC3, 1 / 255.0);
+
+        // Download the final image from GPU to CPU
+        gpuImage.download(resultImage);
+
+        // Create a tensor from the CPU Mat
         torch::Tensor tensorImage = torch::from_blob(
-            image.data, {image.rows, image.cols, 3}, c10::kFloat);
+            resultImage.data, {resultImage.rows, resultImage.cols, 3},
+            torch::kFloat);
         tensorImage = tensorImage.permute({2, 0, 1});
 
-        // Normalize
         std::vector<double> norm_mean = {kImageNormalizationMeanR,
                                          kImageNormalizationMeanG,
                                          kImageNormalizationMeanB};
         std::vector<double> norm_std = {kImageNormalizationStdR,
                                         kImageNormalizationStdG,
                                         kImageNormalizationStdB};
 
+        // Normalize the tensor
         tensorImage = torch::data::transforms::Normalize<>(
             norm_mean, norm_std)(tensorImage);