//
|
// Created by Scheaven on 2020/1/3.
|
//
|
|
#include "reid_feature.h"
|
#include <cuda_runtime_api.h>
|
#include <torch/torch.h>
|
|
bool ReID_Feature::ReID_init(int gpu_id)
|
{
|
if(gpu_id == -1){
|
this->module = torch::jit::load(MODEL_PATH);
|
this->module.to(torch::kCPU);
|
this->module.eval();
|
this->is_gpu = false;
|
}else if(torch::cuda::is_available() && torch::cuda::device_count() >= gpu_id)
|
{
|
cudaSetDevice(gpu_id);
|
cout << "model loading::" << HUMAN_FEATS << endl;
|
this->module = torch::jit::load(MODEL_PATH, torch::Device(torch::DeviceType::CUDA,gpu_id));
|
this->module.to(torch::kCUDA);
|
this->module.eval();
|
this->is_gpu = true;
|
}else{
|
return false;
|
}
|
return true;
|
}
|
|
int ReID_Feature::ReID_size()
|
{
|
int size = 2048;
|
return size;
|
|
}
|
|
bool ReID_Feature::ReID_extractor(float *pBuf, float *pFeature)
|
{
|
auto input_tensor = torch::from_blob(pBuf, {1, 256, 128, 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(this->is_gpu)
|
input_tensor = input_tensor.to(at::kCUDA);
|
torch::Tensor human_feat =this->module.forward({input_tensor}).toTensor();
|
|
// for (int k = 0; k < 20; ++k) {
|
// cout << "--extractor---human_feats------" <<human_feat[0][k+2000]<< endl;
|
// }
|
torch::Tensor query_feat;
|
if(this->is_gpu)
|
query_feat = human_feat.cpu();
|
else
|
query_feat = human_feat;
|
|
auto foo_one = query_feat.accessor<float,2>();
|
|
ReID_Utils RET;
|
|
float f_size = -0.727412;
|
for (int64_t i = 0; i < foo_one.size(0); i++) {
|
auto a1 = foo_one[i];
|
for (int64_t j = 0; j < a1.size(0); j++) {
|
pFeature[j] = a1[j];
|
}
|
}
|
|
// cout << "---- end 11-------" << pFeature[0] << endl;
|
return true;
|
}
|
|
float ReID_Feature::ReID_Compare(float *pFeature1, float *pFeature2)
|
{
|
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] = pFeature1[i];
|
gallery_feat[0][i] = pFeature2[i];
|
}
|
|
if(this->is_gpu)
|
{
|
query_feat = query_feat.cuda();
|
gallery_feat = gallery_feat.cuda();
|
}
|
|
// cout << "-----------------after-----------" << endl;
|
// cout << query_feat<< endl;
|
|
// for (int k = 0; k < 20; ++k) {
|
// cout << "-query_feat----1111111111------" <<query_feat[0][k+2000]<< endl;
|
// }
|
//
|
// cout << "-----------------asdf-----------" << endl;
|
//// cout << gallery_feat[0][0]<< endl;
|
// for (int k = 0; k < 20; ++k) {
|
// cout << "-gallery_feat----22222222------" <<gallery_feat[0][k+2000]<< endl;
|
// }
|
|
torch::Tensor a_similarity = torch::cosine_similarity(query_feat, gallery_feat);
|
|
if(this->is_gpu)
|
a_similarity = a_similarity.cpu();
|
|
auto foo_one = a_similarity.accessor<float,1>();
|
// cout << ":::::::::-" << endl;
|
|
float f_distance = foo_one[0];
|
|
return f_distance;
|
|
}
|
|
void ReID_Feature::ReID_Release()
|
{
|
prinf("release");
|
// this->module = nullptr;//加载模型
|
// return true;
|
}
|
