From 42b3bda03999a16aa9d8be2dca3bb68e90d32b40 Mon Sep 17 00:00:00 2001
From: lancerts <tangshao28@gmail.com>
Date: Mon, 22 Jan 2024 20:12:24 -0800
Subject: [PATCH] Include a CI job for cpp/custom-dataset (#1219)

---
 .github/workflows/main_cpp.yml    | 15 +++++++++++++++
 cpp/custom-dataset/CMakeLists.txt |  5 +++++
 run_cpp_examples.sh               | 32 +++++++++++++++++++++++++++++--
 3 files changed, 50 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/main_cpp.yml b/.github/workflows/main_cpp.yml
index 442cdd7acb..75cf01be11 100644
--- a/.github/workflows/main_cpp.yml
+++ b/.github/workflows/main_cpp.yml
@@ -27,6 +27,21 @@ jobs:
         sudo apt update && sudo apt upgrade
         sudo apt install cmake g++ make
         sudo apt-get -y install intel-mkl
+    - name: Install OpenCV
+      run: |
+        sudo apt -y install libtbb-dev
+        sudo apt install libopencv-dev
+
+    # Alternatively, you can install OpenCV from source
+    # - name: Install OpenCV from source
+    #   run: |
+        # wget -O opencv.zip https://github.com/opencv/opencv/archive/4.x.zip
+        # unzip opencv.zip
+        # mkdir -p build && cd build
+        # cmake  ../opencv-4.x
+        # cmake --build .
+        # sudo make install    
+    
     - name: Run Cpp Tests
       run: |
         chmod +x ./run_cpp_examples.sh
diff --git a/cpp/custom-dataset/CMakeLists.txt b/cpp/custom-dataset/CMakeLists.txt
index fc9d11c8e2..3f45b7946a 100644
--- a/cpp/custom-dataset/CMakeLists.txt
+++ b/cpp/custom-dataset/CMakeLists.txt
@@ -6,6 +6,11 @@ set(CMAKE_CXX_STANDARD 17)
 find_package(Torch REQUIRED)
 find_package(OpenCV REQUIRED COMPONENTS core imgproc imgcodecs)
 
+message(STATUS "OpenCV include dirs: ${OpenCV_INCLUDE_DIRS}")
+message(STATUS "OpenCV libraries: ${OpenCV_LIBS}")
+
+
+include_directories(${OpenCV_INCLUDE_DIRS})
 add_executable(${PROJECT_NAME} "custom-dataset.cpp")
 target_link_libraries(${PROJECT_NAME} "${OpenCV_LIBS}")
 target_link_libraries(${PROJECT_NAME} "${TORCH_LIBRARIES}")
diff --git a/run_cpp_examples.sh b/run_cpp_examples.sh
index 3aed67d023..f0a0ba7802 100644
--- a/run_cpp_examples.sh
+++ b/run_cpp_examples.sh
@@ -22,8 +22,8 @@ function error() {
 function get_libtorch() {
   echo "Getting libtorch"
   cd $HOME_DIR
-  wget https://download.pytorch.org/libtorch/nightly/cpu/libtorch-shared-with-deps-latest.zip
-  unzip libtorch-shared-with-deps-latest.zip
+  wget https://download.pytorch.org/libtorch/nightly/cpu/libtorch-cxx11-abi-shared-with-deps-latest.zip
+  unzip libtorch-cxx11-abi-shared-with-deps-latest.zip
 
   if [ $? -eq 0 ]; then
     echo "Successfully downloaded and extracted libtorch"
@@ -55,6 +55,32 @@ function autograd() {
   fi
 }
 
+function custom-dataset() {
+  start
+  # Download the dataset and unzip it
+  if [ ! -d "$BASE_DIR/cpp/$EXAMPLE/dataset" ]; then
+    wget https://data.caltech.edu/records/mzrjq-6wc02/files/caltech-101.zip
+    unzip caltech-101.zip
+    cd caltech-101
+    tar -xzf 101_ObjectCategories.tar.gz
+    mv 101_ObjectCategories $BASE_DIR/cpp/$EXAMPLE/dataset
+  fi
+  # build the executable and run it
+  cd $BASE_DIR/cpp/$EXAMPLE
+  mkdir build
+  cd build
+  cmake -DCMAKE_PREFIX_PATH=$LIBTORCH_PATH ..
+  make
+  if [ $? -eq 0 ]; then
+    echo "Successfully built $EXAMPLE"
+    cd $BASE_DIR/cpp/$EXAMPLE
+    ./build/$EXAMPLE # Run the executable
+  else
+    error "Failed to build $EXAMPLE"
+    exit 1
+  fi
+}
+
 function clean() {
   cd $BASE_DIR
   echo "Running clean to remove cruft"
@@ -62,11 +88,13 @@ function clean() {
   find . -type d -name 'build' -exec rm -rf {} +
   # Remove the libtorch directory
   rm -rf $HOME_DIR/libtorch
+  rm -f $HOME_DIR/libtorch-shared-with-deps-latest.zip
   echo "Clean completed"
 }
 
 function run_all() {
   autograd
+  custom-dataset
 }
 
 # by default, run all examples