#include "remain_filter.h"
|
#include <cuda_runtime_api.h>
|
#include <torch/torch.h>
|
|
#include <chrono>
|
#include "config_util.h"
|
|
using namespace std;
|
|
RemainFilter* RemainFilter::instance = NULL;
|
|
RemainFilter* RemainFilter::getInstance()
|
{
|
if (instance==NULL)
|
{
|
int gpu_id = 0;
|
instance = new RemainFilter(0);
|
}
|
return instance;
|
}
|
|
RemainFilter::RemainFilter(int gpu_id)
|
{
|
M_STATICSTRUCT config;
|
if (gpu_id==-1)
|
{
|
this->module = torch::jit::load(config.helmet_model);
|
this->module.to(torch::kCPU);
|
this->module.eval();
|
this->is_gpu = false;
|
}else if (torch::cuda::is_available() && torch::cuda::device_count()>= gpu_id)
|
{
|
this->gpu_id = gpu_id;
|
cudaSetDevice(gpu_id);
|
this->module = torch::jit::load(config.helmet_model,torch::Device(torch::DeviceType::CUDA,gpu_id));
|
this->module.to(torch::Device(torch::DeviceType::CUDA, gpu_id));
|
this->module.eval();
|
this->is_gpu = true;
|
}
|
SLOG::getInstance()->addLog(0,"model load ok - gpu_id:" + to_string(gpu_id));
|
}
|
RemainFilter::~RemainFilter(){}
|
|
float RemainFilter::whichRemain(unsigned char *pBuf)
|
{
|
// cv::Mat tmp = cv::imread("/data/disk2/01_dataset/05_ImageNet/val/ILSVRC2012_val_00050000.JPEG");
|
SLOG::getInstance()->addLog(0,"remain_filter.cpp 47: start select abondan: ");
|
auto t1 = std::chrono::steady_clock::now();
|
auto input_tensor = torch::from_blob(pBuf, {1, 224,224,3});
|
input_tensor = input_tensor.permute({0,3,1,2});
|
input_tensor[0][0] = input_tensor[0][0].sub_(0.485).div_(0.229);
|
input_tensor[0][1] = input_tensor[0][1].sub_(0.456).div_(0.224);
|
input_tensor[0][2] = input_tensor[0][2].sub_(0.406).div_(0.225);
|
|
if (is_gpu)
|
{
|
input_tensor = input_tensor.to(torch::Device(torch::DeviceType::CUDA, this->gpu_id));
|
}
|
|
auto out = this->module.forward({input_tensor}).toTensor();
|
auto results = out.sort(-1,true);
|
auto t_len = std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::steady_clock::now() - t1).count();
|
std::cout << "-------FPS----"<<t_len<<std::endl<<std::endl;
|
SLOG::getInstance()->addLog(0,"module use milliseconds:: " + to_string(t_len));
|
|
// std::cout << out<<std::endl;
|
auto pre_value = std::get<0>(results)[0].softmax(0);
|
auto pre_index = std::get<1>(results)[0];
|
|
for (int i = 0; i < 16; ++i) {
|
auto idx = pre_index[i].item<int>();
|
std::cout << " ============= Top-" << i + 1
|
<< " =============" << std::endl;
|
std::cout << " Label: " << idx << std::endl;
|
std::cout << " With Probability: "
|
<< pre_value[i].item<float>() * 100.0f << "%" << std::endl;
|
}
|
|
if(pre_index[0].item<int>()<5)
|
{
|
SLOG::getInstance()->addLog(0,"abandon is human!! results ID is: " + to_string(pre_index[0].item<int>()));
|
return 0.0f;
|
}
|
|
return pre_value[0].item<float>() * 100.0f;
|
|
// tmp;
|
// std::cout << std::get<0>(results)[0].softmax(0) << "----------out::"<<std::endl;
|
|
}
|