| |
| |
|
|
|
|
| from cv2box import CVImage, MyFpsCounter |
|
|
| from model_lib import ModelBase |
| import numpy as np |
| import cv2 |
|
|
| MODEL_ZOO = { |
| 'xseg_0611': { |
| 'model_path': './pretrain_models/xseg_230611_16_17.onnx', |
| 'input_dynamic_shape': [[1, 256, 256, 3]] |
| }, |
| } |
|
|
|
|
| class XSEG(ModelBase): |
| def __init__(self, model_type='xseg_0611', provider='cpu'): |
| super().__init__(MODEL_ZOO[model_type], provider) |
| self.model_type = model_type |
|
|
|
|
| def forward(self, face_image): |
| """ |
| Args: |
| face_image: cv2 image -1~1 RGB |
| Returns: |
| RGB 256x256x3 -1~1 |
| """ |
| face_image = (face_image + 1) / 2 |
| if face_image.shape[-1] >= 4: |
| if len(face_image.shape)>3: |
| face_image = face_image[0] |
| face_image = face_image.transpose(1, 2, 0) |
| face_image_h, face_image_w, _ = face_image.shape |
| if face_image_h != 256: |
| face_image = cv2.resize(face_image, (256, 256)) |
| image_out = self.model.forward(face_image[...,::-1][None].astype(np.float32)) |
| |
| output_face = (image_out[0].squeeze()).clip(0, 1) |
| if face_image_h != 256: |
| output_face = cv2.resize(output_face, (face_image_w, face_image_h)) |
| return output_face |
|
|
|
|
| if __name__ == '__main__': |
| face_img_p = 'data/source/ym-1.jpeg' |
| fa = XSEG(model_type='xseg_0611', provider='trt16') |
| face_img = (cv2.resize(cv2.imread(face_img_p)/127.5-1,(512,512)))[...,::-1] |
|
|
| with MyFpsCounter() as mfc: |
| for i in range(20): |
| face = fa.forward(face_img) |
| |
| |
|
|
