1+ #include " Face_Alignment.h"
2+
3+ Face_Alignment::Face_Alignment ()
4+ {
5+ }
6+
7+ Face_Alignment::~Face_Alignment ()
8+ {
9+ }
10+ HZFLAG Face_Alignment::Face_AlignmentInit (Config&config)
11+ {
12+ Face_Alignment_INPUT_BLOB_NAME = " data" ;
13+ Face_Alignment_OUTPUT_BLOB_NAME = " fc1" ;
14+ Face_Alignment_INPUT_H = 192 ;
15+ Face_Alignment_INPUT_W = 192 ;
16+ Face_Alignment_OUTPUT_SIZE = 212 ;
17+ cudaSetDevice (config.gpu_id );
18+ char *trtModelStream{nullptr };
19+ size_t size{0 };
20+ std::string directory;
21+ const size_t last_slash_idx = config.FaceAlignmentModelPath .rfind (" .onnx" );
22+ if (std::string::npos != last_slash_idx)
23+ {
24+ directory = config.FaceAlignmentModelPath .substr (0 , last_slash_idx);
25+ }
26+ std::string out_engine = directory +" _batch=" + std::to_string (config.FaceAlignment_bs ) + " .engine" ;
27+ bool enginemodel = Face_Alignment_model_exists (out_engine);
28+ if (!enginemodel)
29+ {
30+ std::cout << " Building engine, please wait for a while..." << std::endl;
31+ bool wts_model = Face_Alignment_model_exists (config.FaceAlignmentModelPath );
32+ if (!wts_model)
33+ {
34+ std::cout << " ONNX file is not Exist!Please Check!" << std::endl;
35+ return HZ_ERROR;
36+ }
37+ Onnx2Ttr onnx2trt;
38+ onnx2trt.onnxToTRTModel (Face_Alignment_gLogger,config.FaceAlignmentModelPath .c_str (),config.FaceAlignment_bs ,out_engine.c_str ());// config.classs_path
39+ }
40+ std::ifstream file (out_engine, std::ios::binary);
41+ if (file.good ())
42+ {
43+ file.seekg (0 , file.end );
44+ size = file.tellg ();
45+ file.seekg (0 , file.beg );
46+ trtModelStream = new char [size];
47+ assert (trtModelStream);
48+ file.read (trtModelStream, size);
49+ file.close ();
50+ }
51+ Face_Alignment_runtime = createInferRuntime (Face_Alignment_gLogger);
52+ assert (Face_Alignment_runtime != nullptr );
53+ Face_Alignment_engine = Face_Alignment_runtime->deserializeCudaEngine (trtModelStream, size);
54+ assert (Face_Alignment_engine != nullptr );
55+ Face_Alignment_context =Face_Alignment_engine->createExecutionContext ();
56+ assert (Face_Alignment_context != nullptr );
57+ delete[] trtModelStream;
58+
59+ Face_Alignment_data = new float [3 * Face_Alignment_INPUT_H * Face_Alignment_INPUT_W];
60+ Face_Alignment_prob = new float [Face_Alignment_OUTPUT_SIZE];
61+
62+ return HZ_SUCCESS;
63+ }
64+ HZFLAG Face_Alignment::Face_AlignmentRun (cv::Mat&img, AlignmentFace&alignmentface)
65+ {
66+ if (img.empty ())
67+ {
68+ std::cout<<" Face_AlignmentRun image is empty" <<std::endl;
69+ return HZ_IMGEMPTY;
70+ }
71+ cv::Mat pr_img;
72+ cv::resize (img,pr_img, cv::Size (Face_Alignment_INPUT_H, Face_Alignment_INPUT_W));
73+ int i = 0 ;
74+ for (int row = 0 ; row < Face_Alignment_INPUT_H; ++row)
75+ {
76+ uchar* uc_pixel = pr_img.data + row * pr_img.step ;
77+ for (int col = 0 ; col < Face_Alignment_INPUT_W; ++col)
78+ {
79+ Face_Alignment_data[i] = (float )uc_pixel[2 ];
80+ Face_Alignment_data[i + Face_Alignment_INPUT_H * Face_Alignment_INPUT_W] = (float )uc_pixel[1 ];
81+ Face_Alignment_data[i + 2 * Face_Alignment_INPUT_H * Face_Alignment_INPUT_W] = (float )uc_pixel[0 ];
82+ uc_pixel += 3 ;
83+ ++i;
84+ }
85+ }
86+ // Run inference
87+ auto start = std::chrono::system_clock::now ();
88+ Face_Alignment_doInference (*Face_Alignment_context,Face_Alignment_data,Face_Alignment_prob, 1 );
89+ auto end = std::chrono::system_clock::now ();
90+ for (int start = 0 ; start < 212 ; start += 2 )
91+ {
92+ alignmentface.landmarks [start]=int ((Face_Alignment_prob[start]+1 )*img.cols / 2 );
93+ alignmentface.landmarks [start+1 ]=int ((Face_Alignment_prob[start+1 ]+1 )*img.cols / 2 );
94+ // cv::circle(img, cv::Point(alignmentface.landmarks[start], alignmentface.landmarks[start+1]), 1, cv::Scalar(200, 160, 75), -1, cv::LINE_8, 0);
95+ }
96+ // 106 keypoint
97+ return HZ_SUCCESS;
98+ }
99+ HZFLAG Face_Alignment::Face_AlignmentRelease ()
100+ {
101+ Face_Alignment_context->destroy ();
102+ Face_Alignment_engine->destroy ();
103+ Face_Alignment_runtime->destroy ();
104+ delete[] Face_Alignment_data;
105+ delete[] Face_Alignment_prob;
106+ Face_Alignment_data = NULL ;
107+ Face_Alignment_prob = NULL ;
108+ return HZ_SUCCESS;
109+ }
110+ void Face_Alignment::Face_Alignment_doInference (IExecutionContext& context, float * input, float * output, int batchSize)
111+ {
112+ const ICudaEngine& engine = context.getEngine ();
113+
114+ // Pointers to input and output device buffers to pass to engine.
115+ // Engine requires exactly IEngine::getNbBindings() number of buffers.
116+ assert (engine.getNbBindings () == 2 );
117+ void * buffers[2 ];
118+
119+ // In order to bind the buffers, we need to know the names of the input and output tensors.
120+ // Note that indices are guaranteed to be less than IEngine::getNbBindings()
121+ const int inputIndex = engine.getBindingIndex (Face_Alignment_INPUT_BLOB_NAME);
122+ const int outputIndex = engine.getBindingIndex (Face_Alignment_OUTPUT_BLOB_NAME);
123+
124+ // Create GPU buffers on device
125+ CHECK (cudaMalloc (&buffers[inputIndex], batchSize * 3 * Face_Alignment_INPUT_H * Face_Alignment_INPUT_W * sizeof (float )));
126+ CHECK (cudaMalloc (&buffers[outputIndex], batchSize * Face_Alignment_OUTPUT_SIZE * sizeof (float )));
127+
128+ // Create stream
129+ cudaStream_t stream;
130+ CHECK (cudaStreamCreate (&stream));
131+
132+ // DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
133+ CHECK (cudaMemcpyAsync (buffers[inputIndex], input, batchSize * 3 * Face_Alignment_INPUT_H * Face_Alignment_INPUT_W * sizeof (float ), cudaMemcpyHostToDevice, stream));
134+ context.enqueue (batchSize, buffers, stream, nullptr );
135+ CHECK (cudaMemcpyAsync (output, buffers[outputIndex], batchSize * Face_Alignment_OUTPUT_SIZE * sizeof (float ), cudaMemcpyDeviceToHost, stream));
136+ cudaStreamSynchronize (stream);
137+
138+ // Release stream and buffers
139+ cudaStreamDestroy (stream);
140+ CHECK (cudaFree (buffers[inputIndex]));
141+ CHECK (cudaFree (buffers[outputIndex]));
142+ }
143+ bool Face_Alignment::Face_Alignment_model_exists (const std::string& name)
144+ {
145+ std::ifstream f (name.c_str ());
146+ return f.good ();
147+ }
0 commit comments