Yaxin
/
Wafer_Inspection
31
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

优化代码

Browse Source
main
熊继淙 1 year ago
parent 4be8422fda
commit e1119953ac
2 changed files with 22 additions and 22 deletions
Show Stats Download Patch File Download Diff File

42
Code/MF_TRTInfer/utils.cpp
Unescape Escape View File

@ -1,6 +1,6 @@
#include "utils.h" #include "utils.h"
void utils::saveBinaryFile(float* vec, size_t len, const std::string& file) void trtUtils::saveBinaryFile(float* vec, size_t len, const std::string& file)
{ {
std::ofstream out(file, std::ios::out | std::ios::binary); std::ofstream out(file, std::ios::out | std::ios::binary);
if (!out.is_open()) if (!out.is_open())
@ -9,7 +9,7 @@ void utils::saveBinaryFile(float* vec, size_t len, const std::string& file)
out.close(); out.close();
} }
std::vector<uint8_t> utils::readBinaryFile(const std::string& file) std::vector<uint8_t> trtUtils::readBinaryFile(const std::string& file)
{ {
std::ifstream in(file, std::ios::in | std::ios::binary); std::ifstream in(file, std::ios::in | std::ios::binary);
if (!in.is_open()) if (!in.is_open())
@ -30,7 +30,7 @@ std::vector<uint8_t> utils::readBinaryFile(const std::string& file)
} }
std::vector<unsigned char> utils::loadModel(const std::string& file) std::vector<unsigned char> trtUtils::loadModel(const std::string& file)
{ {
std::ifstream in(file, std::ios::in | std::ios::binary); std::ifstream in(file, std::ios::in | std::ios::binary);
if (!in.is_open()) if (!in.is_open())
@ -51,19 +51,19 @@ std::vector<unsigned char> utils::loadModel(const std::string& file)
return data; return data;
} }
std::string utils::getSystemTimeStr() std::string trtUtils::getSystemTimeStr()
{ {
return std::to_string(std::rand()); return std::to_string(std::rand());
} }
bool utils::setInputStream(const utils::InputStream& source, const std::string& imagePath, const std::string& videoPath, const int& cameraID, bool trtUtils::setInputStream(const trtUtils::InputStream& source, const std::string& imagePath, const std::string& videoPath, const int& cameraID,
cv::VideoCapture& capture, int& totalBatches, int& delayTime, utils::InitParameter& param) cv::VideoCapture& capture, int& totalBatches, int& delayTime, trtUtils::InitParameter& param)
{ {
int total_frames = 0; int total_frames = 0;
std::string img_format; std::string img_format;
switch (source) switch (source)
{ {
case utils::InputStream::IMAGE: case trtUtils::InputStream::IMAGE:
img_format = imagePath.substr(imagePath.size()-4, 4); img_format = imagePath.substr(imagePath.size()-4, 4);
if (img_format == ".png" || img_format == ".PNG") if (img_format == ".png" || img_format == ".PNG")
{ {
@ -79,13 +79,13 @@ bool utils::setInputStream(const utils::InputStream& source, const std::string&
totalBatches = 1; totalBatches = 1;
delayTime = 0; delayTime = 0;
break; break;
case utils::InputStream::VIDEO: case trtUtils::InputStream::VIDEO:
capture.open(videoPath); capture.open(videoPath);
total_frames = capture.get(cv::CAP_PROP_FRAME_COUNT); total_frames = capture.get(cv::CAP_PROP_FRAME_COUNT);
totalBatches = (total_frames % param.batch_size == 0) ? totalBatches = (total_frames % param.batch_size == 0) ?
(total_frames / param.batch_size) : (total_frames / param.batch_size + 1); (total_frames / param.batch_size) : (total_frames / param.batch_size + 1);
break; break;
case utils::InputStream::CAMERA: case trtUtils::InputStream::CAMERA:
capture.open(cameraID); capture.open(cameraID);
total_frames = INT_MAX; total_frames = INT_MAX;
totalBatches = INT_MAX; totalBatches = INT_MAX;
@ -99,7 +99,7 @@ bool utils::setInputStream(const utils::InputStream& source, const std::string&
return capture.isOpened(); return capture.isOpened();
} }
void utils::setRenderWindow(InitParameter& param) void trtUtils::setRenderWindow(InitParameter& param)
{ {
if (!param.is_show) if (!param.is_show)
return; return;
@ -122,14 +122,14 @@ void utils::setRenderWindow(InitParameter& param)
} }
} }
std::string utils::getTimeStamp() std::string trtUtils::getTimeStamp()
{ {
std::chrono::nanoseconds t = std::chrono::duration_cast<std::chrono::nanoseconds>( std::chrono::nanoseconds t = std::chrono::duration_cast<std::chrono::nanoseconds>(
std::chrono::system_clock::now().time_since_epoch()); std::chrono::system_clock::now().time_since_epoch());
return std::to_string(t.count()); return std::to_string(t.count());
} }
void utils::show(const std::vector<std::vector<utils::Box>>& objectss, const std::vector<std::string>& classNames, void trtUtils::show(const std::vector<std::vector<trtUtils::Box>>& objectss, const std::vector<std::string>& classNames,
const int& cvDelayTime, std::vector<cv::Mat>& imgsBatch) const int& cvDelayTime, std::vector<cv::Mat>& imgsBatch)
{ {
std::string windows_title = "image"; std::string windows_title = "image";
@ -191,7 +191,7 @@ void utils::show(const std::vector<std::vector<utils::Box>>& objectss, const std
} }
} }
void utils::save(const std::vector<std::vector<Box>>& objectss, const std::vector<std::string>& classNames, void trtUtils::save(const std::vector<std::vector<trtUtils::Box>>& objectss, const std::vector<std::string>& classNames,
const std::string& savePath, std::vector<cv::Mat>& imgsBatch, const int& batchSize, const int& batchi) const std::string& savePath, std::vector<cv::Mat>& imgsBatch, const int& batchSize, const int& batchi)
{ {
cv::Scalar color = cv::Scalar(0, 255, 0); cv::Scalar color = cv::Scalar(0, 255, 0);
@ -241,30 +241,30 @@ void utils::save(const std::vector<std::vector<Box>>& objectss, const std::vecto
} }
} }
utils::HostTimer::HostTimer() trtUtils::HostTimer::HostTimer()
{ {
t1 = std::chrono::steady_clock::now(); t1 = std::chrono::steady_clock::now();
} }
float utils::HostTimer::getUsedTime() float trtUtils::HostTimer::getUsedTime()
{ {
t2 = std::chrono::steady_clock::now(); t2 = std::chrono::steady_clock::now();
std::chrono::duration<double> time_used = std::chrono::duration_cast<std::chrono::duration<double>>(t2 - t1); std::chrono::duration<double> time_used = std::chrono::duration_cast<std::chrono::duration<double>>(t2 - t1);
return(1000 * time_used.count()); // ms return(1000 * time_used.count()); // ms
} }
utils::HostTimer::~HostTimer() trtUtils::HostTimer::~HostTimer()
{ {
} }
utils::DeviceTimer::DeviceTimer() trtUtils::DeviceTimer::DeviceTimer()
{ {
cudaEventCreate(&start); cudaEventCreate(&start);
cudaEventCreate(&end); cudaEventCreate(&end);
cudaEventRecord(start); cudaEventRecord(start);
} }
float utils::DeviceTimer::getUsedTime() float trtUtils::DeviceTimer::getUsedTime()
{ {
cudaEventRecord(end); cudaEventRecord(end);
cudaEventSynchronize(end); cudaEventSynchronize(end);
@ -273,14 +273,14 @@ float utils::DeviceTimer::getUsedTime()
return total_time; return total_time;
} }
utils::DeviceTimer::DeviceTimer(cudaStream_t stream) trtUtils::DeviceTimer::DeviceTimer(cudaStream_t stream)
{ {
cudaEventCreate(&start); cudaEventCreate(&start);
cudaEventCreate(&end); cudaEventCreate(&end);
cudaEventRecord(start, stream); cudaEventRecord(start, stream);
} }
float utils::DeviceTimer::getUsedTime(cudaStream_t stream) float trtUtils::DeviceTimer::getUsedTime(cudaStream_t stream)
{ {
cudaEventRecord(end, stream); cudaEventRecord(end, stream);
cudaEventSynchronize(end); cudaEventSynchronize(end);
@ -289,7 +289,7 @@ float utils::DeviceTimer::getUsedTime(cudaStream_t stream)
return total_time; return total_time;
} }
utils::DeviceTimer::~DeviceTimer() trtUtils::DeviceTimer::~DeviceTimer()
{ {
cudaEventDestroy(start); cudaEventDestroy(start);
cudaEventDestroy(end); cudaEventDestroy(end);

2
Code/MF_TRTInfer/utils.h
Unescape Escape View File

@ -14,7 +14,7 @@
#endif #endif
namespace utils namespace trtUtils
{ {
namespace dataSets namespace dataSets
{ {

Write Preview
Loading…
Cancel
Save