64 lines
1.9 KiB
C++
Executable File
64 lines
1.9 KiB
C++
Executable File
#ifndef __TRT_UTILS_H_
|
|
#define __TRT_UTILS_H_
|
|
|
|
#include <iostream>
|
|
#include <vector>
|
|
#include <algorithm>
|
|
#include <cudnn.h>
|
|
#include <dirent.h>
|
|
#include <opencv2/opencv.hpp>
|
|
|
|
#ifndef CUDA_CHECK
|
|
|
|
#define CUDA_CHECK(callstr) \
|
|
{ \
|
|
cudaError_t error_code = callstr; \
|
|
if (error_code != cudaSuccess) { \
|
|
std::cerr << "CUDA error " << error_code << " at " << __FILE__ << ":" << __LINE__; \
|
|
assert(0); \
|
|
} \
|
|
}
|
|
|
|
#endif
|
|
|
|
namespace Tn
|
|
{
|
|
template<typename T>
|
|
void write(char*& buffer, const T& val)
|
|
{
|
|
*reinterpret_cast<T*>(buffer) = val;
|
|
buffer += sizeof(T);
|
|
}
|
|
|
|
template<typename T>
|
|
void read(const char*& buffer, T& val)
|
|
{
|
|
val = *reinterpret_cast<const T*>(buffer);
|
|
buffer += sizeof(T);
|
|
}
|
|
}
|
|
|
|
static inline int read_files_in_dir(const char *p_dir_name, std::vector<std::string> &file_names) {
|
|
DIR *p_dir = opendir(p_dir_name);
|
|
if (p_dir == nullptr) {
|
|
return -1;
|
|
}
|
|
|
|
struct dirent* p_file = nullptr;
|
|
while ((p_file = readdir(p_dir)) != nullptr) {
|
|
if (strcmp(p_file->d_name, ".") != 0 &&
|
|
strcmp(p_file->d_name, "..") != 0) {
|
|
//std::string cur_file_name(p_dir_name);
|
|
//cur_file_name += "/";
|
|
//cur_file_name += p_file->d_name;
|
|
std::string cur_file_name(p_file->d_name);
|
|
file_names.push_back(cur_file_name);
|
|
}
|
|
}
|
|
|
|
closedir(p_dir);
|
|
return 0;
|
|
}
|
|
|
|
#endif
|