diff --git a/.gitignore b/.gitignore index 534fed2..4bdf96e 100644 --- a/.gitignore +++ b/.gitignore @@ -109,3 +109,6 @@ venv.bak/ *.jpeg *.pt .DS_Store +src/libtorch/ +*.sh +.vscode/settings.json diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt new file mode 100644 index 0000000..dba2550 --- /dev/null +++ b/src/CMakeLists.txt @@ -0,0 +1,10 @@ +cmake_minimum_required(VERSION 3.0 FATAL_ERROR) +project(custom_ops) +# set(OpenCV_DIR /home/zjs/opt/OpenCV-3.4.0-cuda-9.0/share/OpenCV) +# include (CTest) +find_package(OpenCV REQUIRED) +include_directories(${OpenCV_INCLUDE_DIRS}) +find_package(Torch REQUIRED) +add_executable(main main.cpp) +target_link_libraries(main ${TORCH_LIBRARIES} ${OpenCV_LIBS}) +set_property(TARGET main PROPERTY CXX_STANDARD 14) \ No newline at end of file diff --git a/src/main.cpp b/src/main.cpp new file mode 100644 index 0000000..555ead8 --- /dev/null +++ b/src/main.cpp @@ -0,0 +1,117 @@ +#include // One-stop header. +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#define GPU 1 + +template +T remove_extension(T const & filename) +{ + typename T::size_type const p(filename.find_last_of('.')); + return p > 0 && p != T::npos ? filename.substr(0, p) : filename; +} +std::string base_name(std::string const & path) +{ + return path.substr(path.find_last_of("/\\") + 1); +} +int main(int argc, const char* argv[]) { + if (argc != 3) { + std::cerr << "usage: main \n"; + return -1; + } + + torch::jit::script::Module module; + try { + // Deserialize the ScriptModule from a file using torch::jit::load(). + module = torch::jit::load(argv[1]); + std::cout<<"Module loaded successfuly"< fn; + glob(argv[2], fn, false); + + std::vector images; + size_t count = fn.size(); //number of png files in images folder + for (size_t i=0; i({row, col})).align_corners(false)); + + net_out = net_out.squeeze(); + net_out = (net_out.min(net_out.max()) / (net_out.max() - net_out.min())).mul(65535.0f); + net_out = net_out.to(torch::kCPU); + int height = net_out.sizes()[0]; + int width = net_out.sizes()[1]; + int x = net_out.sizes()[2]; + + try + { + cv::Mat output_mat(cv::Size{width, height}, CV_32FC1, net_out.data_ptr()); + output_mat.convertTo(output_mat, CV_16UC1); + double fps = cv::getTickFrequency() / (cv::getTickCount() - start); + std::cout << "FPS : " << fps << std::endl; + std::string output_address = "./output/"; + + output_address.append(remove_extension(base_name(fn[i])).append(extention)); + cv::imwrite(output_address, output_mat); + + } + catch (const c10::Error& e) + { + std::cout << "an error has occured : " << e.msg() << std::endl; + } + } + +} +