#include "reid.h"
|
|
#include <stdio.h>
|
#include <stdexcept>
|
// #include <cuda_runtime_api.h>
|
#include <torch/torch.h>
|
#include "opencv2/opencv.hpp"
|
|
reid::reid(const int gpu_index, const char *module_path)
|
:is_gpu_(gpu_index >= 0)
|
,start_(false)
|
{
|
if(init(gpu_index, module_path) < 0){
|
throw std::runtime_error("init swscale error");
|
}
|
}
|
|
reid::~reid()
|
{
|
release();
|
}
|
|
int reid::init(const int gpu_index, const char *module_path){
|
const int gpu_id = gpu_index;
|
if(gpu_id == -1){
|
module_ = torch::jit::load(module_path);
|
module_.to(torch::kCPU);
|
module_.eval();
|
is_gpu_ = false;
|
printf("==========load CPU\n");
|
}else if(torch::cuda::is_available() && torch::cuda::device_count() >= gpu_id){
|
// cudaSetDevice(gpu_id);
|
module_ = torch::jit::load(module_path, torch::Device(torch::DeviceType::CUDA, gpu_id));
|
module_.to(torch::kCUDA);
|
module_.eval();
|
is_gpu_ = true;
|
printf("==========load GPU\n");
|
}else{
|
printf("reid use gpu %d, error\n", gpu_id);
|
return -1;
|
}
|
start_ = true;
|
return 0;
|
}
|
|
void reid::release(){
|
|
}
|
|
int reid::extract(unsigned char *img, const int w, const int h, const int c, float *feature){
|
if (!start_ || !img) return -1;
|
|
printf("&&&&&& Start ReID Image Size: %dx%d\n", w, h);
|
printf("\n\n\n\n=============================================");
|
|
cv::Mat matImg(h, w, CV_8UC3, img);
|
|
//cv::VideoCapture capture2;
|
//cv::Mat matImg;
|
//matImg = capture2.open("rtsp://admin:a1234567@192.168.5.34:554/h264/ch1/main/av_stream");
|
//capture2.read(matImg);
|
|
|
|
// scale image to fit
|
cv::Size scale(256, 128);
|
cv::resize(matImg, matImg, scale);
|
|
// vector<int> compression_params;
|
// compression_params.push_back(cv::IMWRITE_PNG_COMPRESSION);
|
// compression_params.push_back(9);
|
|
//cv::imwrite("0000.png", matImg, compression_params);
|
matImg.convertTo(matImg, CV_32FC3, 1.0f/255.0f);
|
printf("&&&&&& matImg.convertTo\n");
|
|
auto input_tensor = torch::from_blob(matImg.data, {1, 256, 128, 3});
|
printf("&&&&&& torch::from_blob\n");
|
|
input_tensor = input_tensor.permute({0, 3, 1, 2});
|
printf("&&&&&& input_tensor.permute\n");
|
|
input_tensor[0][0] = input_tensor[0][0].sub_(0.485).div_(0.229);
|
printf("&&&&&& input_tensor[0][0]\n");
|
|
input_tensor[0][1] = input_tensor[0][1].sub_(0.456).div_(0.224);
|
printf("&&&&&& input_tensor[0][1]\n");
|
|
input_tensor[0][2] = input_tensor[0][2].sub_(0.406).div_(0.225);
|
printf("&&&&&& input_tensor[0][2]\n");
|
|
if(is_gpu_)
|
{
|
input_tensor = input_tensor.to(at::kCUDA);
|
printf("&&&&&& kCUDA\n");
|
}else
|
{
|
printf("------ kCPU)\n");
|
}
|
|
|
torch::Tensor human_feat = module_.forward({input_tensor}).toTensor();
|
printf("------human_feat---\n");
|
torch::Tensor query_feat;
|
|
if(is_gpu_){
|
printf("------is_gpu_--\n");
|
query_feat = human_feat.cpu();
|
}
|
else
|
query_feat = human_feat;
|
|
|
printf("&&&&&& query_feat\n");
|
|
auto foo_one = query_feat.accessor<float,2>();
|
printf("&&&&&& query_feat.accessor\n");
|
|
float f_size = -0.727412;
|
for (int64_t i = 0; i < foo_one.size(0); i++) {
|
printf("------foo_one.size(0)::%d\n",foo_one.size(0));
|
auto a1 = foo_one[i];
|
for (int64_t j = 0; j < a1.size(0); j++) {
|
feature[j] = a1[j];
|
}
|
}
|
|
// for (int k = 0; k < 3; ++k) {
|
// printf("--extractor---human_feats------%f",feature[k+2000]);
|
// }
|
|
printf("\n\n\n\n\n----------------------------------------------------------------------------\n");
|
|
printf("&&&&&& End ReID With CUDA %d\n", is_gpu_);
|
|
return 0;
|
}
|
|
float reid::compare(const float *feat1, const float *feat2){
|
torch::Tensor query_feat = torch::zeros({1,2048});
|
torch::Tensor gallery_feat = torch::zeros({1,2048});
|
|
for (int i = 0; i < 2048; i++) {
|
query_feat[0][i] = feat1[i];
|
gallery_feat[0][i] = feat2[i];
|
}
|
|
if(is_gpu_) {
|
query_feat = query_feat.cuda();
|
gallery_feat = gallery_feat.cuda();
|
}
|
|
torch::Tensor a_similarity = torch::cosine_similarity(query_feat, gallery_feat);
|
|
if(is_gpu_)
|
a_similarity = a_similarity.cpu();
|
|
auto foo_one = a_similarity.accessor<float,1>();
|
|
float f_distance = foo_one[0];
|
|
return f_distance;
|
}
|
