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