From b35529cf41aa7b75f2e7bac90f777e390b679a8a Mon Sep 17 00:00:00 2001
From: mht <tohidihassan9@gmail.com>
Date: Wed, 28 May 2025 22:53:40 +0330
Subject: [PATCH] Fix: CUDA setup, demo script errors, weight paths &
 .gitignore

---
 .gitignore                         | 39 ++++++++++++++++++++++++++++++
 build.sh                           |  3 +++
 cimp/bb_regressor/bb_regressor.cpp |  3 +++
 cimp/demo.cpp                      |  8 ++++--
 run_demo.sh                        |  6 +++++
 test/run_tests.sh                  | 37 ++++++++++++++++++++++++++++
 6 files changed, 94 insertions(+), 2 deletions(-)
 create mode 100644 .gitignore
 create mode 100644 test/run_tests.sh

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..b2dbcb0
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,39 @@
+# Build artifacts
+build/
+cmake-build-*/
+*.o
+*.a
+*.so
+*.dll
+*.exe
+*.out
+
+# Binaries
+bin/
+
+# Python cache
+__pycache__/
+*.pyc
+*.pyd
+
+# IDE specific
+.idea/
+*.suo
+*.ntvs*
+*.njsproj
+*.sln
+*.sw?
+
+# Generated data/log files
+*.log
+*.txt
+!README.txt
+
+# Test outputs
+test/output/
+test/input_samples/
+test/comparison/
+
+# OS specific
+.DS_Store
+Thumbs.db 
\ No newline at end of file
diff --git a/build.sh b/build.sh
index a2c5463..4f17874 100644
--- a/build.sh
+++ b/build.sh
@@ -1,5 +1,8 @@
 #!/bin/bash
 
+export PATH=/usr/local/cuda-11.5/bin:$PATH
+export LD_LIBRARY_PATH=/usr/local/cuda-11.5/lib64:$LD_LIBRARY_PATH
+
 # Exit on error
 set -e
 
diff --git a/cimp/bb_regressor/bb_regressor.cpp b/cimp/bb_regressor/bb_regressor.cpp
index 900f413..67fed25 100644
--- a/cimp/bb_regressor/bb_regressor.cpp
+++ b/cimp/bb_regressor/bb_regressor.cpp
@@ -291,6 +291,9 @@ BBRegressor::BBRegressor(const std::string& model_weights_dir, torch::Device dev
     // Set the model to evaluation mode
     this->eval();
     
+    // Move the model to the specified device
+    this->to(device);
+    
     // Debug information
     std::cout << "BB Regressor initialized in evaluation mode" << std::endl;
 }
diff --git a/cimp/demo.cpp b/cimp/demo.cpp
index c3b1e41..8007d08 100644
--- a/cimp/demo.cpp
+++ b/cimp/demo.cpp
@@ -167,14 +167,18 @@ int main(int argc, char* argv[]) {
         
         std::cout << "Using exported weights from: " << std::filesystem::absolute(base_dir).string() << std::endl << std::endl;
         
+        std::filesystem::path actual_weights_dir_path = std::filesystem::path(base_dir) / "exported_weights";
+        std::string bb_regressor_weights_dir = (actual_weights_dir_path / "bb_regressor").string();
+        std::string classifier_weights_dir = (actual_weights_dir_path / "classifier").string();
+        
         // Initialize BBRegressor and Classifier
         std::cout << "Initializing BBRegressor..." << std::endl;
-        BBRegressor bb_regressor(base_dir, device);
+        BBRegressor bb_regressor(bb_regressor_weights_dir, device);
         bb_regressor.print_model_info();
         std::cout << std::endl;
         
         std::cout << "Initializing Classifier..." << std::endl;
-        Classifier classifier(base_dir, device);
+        Classifier classifier(classifier_weights_dir, device);
         classifier.print_model_info();
         std::cout << std::endl;
         
diff --git a/run_demo.sh b/run_demo.sh
index 2b56264..f4f3ea0 100644
--- a/run_demo.sh
+++ b/run_demo.sh
@@ -1,5 +1,11 @@
 #!/bin/bash
 
+export PATH=/usr/local/cuda-11.5/bin:$PATH
+export LD_LIBRARY_PATH=/usr/local/cuda-11.5/lib64:$LD_LIBRARY_PATH
+
+# Exit on error
+# set -e # Commented out to allow checking specific error codes if needed
+
 # Set CUDA environment
 if [ -d "/usr/local/cuda" ]; then
     export CUDA_HOME=/usr/local/cuda
diff --git a/test/run_tests.sh b/test/run_tests.sh
new file mode 100644
index 0000000..09e2aeb
--- /dev/null
+++ b/test/run_tests.sh
@@ -0,0 +1,37 @@
+#!/bin/bash
+
+export PATH=/usr/local/cuda-11.5/bin:$PATH
+export LD_LIBRARY_PATH=/usr/local/cuda-11.5/lib64:$LD_LIBRARY_PATH
+
+# Set the number of test samples (default: 1000, but can be reduced for faster testing)
+NUM_SAMPLES=100
+if [ "$1" != "" ]; then
+    NUM_SAMPLES=$1
+fi
+
+echo "Running tests with $NUM_SAMPLES samples..."
+
+# Get the directory of this script
+SCRIPT_DIR=$(dirname "$(readlink -f "$0")")
+PROJECT_DIR=$(dirname "$SCRIPT_DIR")
+cd "$PROJECT_DIR"
+
+# Make sure the build directory exists
+mkdir -p build
+
+# Build the C++ code
+cd build
+echo "Building C++ code..."
+cmake ..
+make -j$(nproc)
+
+# Run the C++ test program to generate samples and outputs
+echo "Running C++ tests..."
+./test_models ../test/input_samples ../test/output $NUM_SAMPLES
+
+# Run the Python comparison script
+cd "$PROJECT_DIR"
+echo "Running Python comparison..."
+python test/compare_models.py --num-samples $NUM_SAMPLES
+
+echo "Tests complete! View results in test/comparison/report.html" 
\ No newline at end of file