| 305 | |
| 306 | |
| 307 | class FaceDet(NanoDetABC): |
| 308 | def __init__(self, model_path="", gpu_id=None, *args, **kwargs): |
| 309 | super(FaceDet, self).__init__(*args, **kwargs) |
| 310 | |
| 311 | self.model_path = model_path |
| 312 | self.ort_session = create_onnx_session(model_path, gpu_id=gpu_id) |
| 313 | self.input_name = self.ort_session.get_inputs()[0].name |
| 314 | |
| 315 | def infer_image(self, img_input): |
| 316 | inference_results = self.ort_session.run(None, {self.input_name: img_input}) |
| 317 | |
| 318 | scores = [np.squeeze(x) for x in inference_results[:3]] |
| 319 | raw_boxes = [np.squeeze(x) for x in inference_results[3:]] |
| 320 | return scores, raw_boxes |