Shellbrady/LivePortrait5

1# coding: utf-82 3import os.path as osp4import cv2; cv2.setNumThreads(0); cv2.ocl.setUseOpenCL(False)5import torch6import numpy as np7import onnxruntime8from .timer import Timer9from .rprint import rlog10from .crop import crop_image, _transform_pts11 12 13def make_abs_path(fn):14    return osp.join(osp.dirname(osp.realpath(__file__)), fn)15 16 17def to_ndarray(obj):18    if isinstance(obj, torch.Tensor):19        return obj.cpu().numpy()20    elif isinstance(obj, np.ndarray):21        return obj22    else:23        return np.array(obj)24 25 26class LandmarkRunner(object):27    """landmark runner"""28    def __init__(self, **kwargs):29        ckpt_path = kwargs.get('ckpt_path')30        onnx_provider = 'cpu' #run on cpu for it to work with ZeroGPU // kwargs.get('onnx_provider', 'cuda')  # 默认用cuda31        device_id = kwargs.get('device_id', 0)32        self.dsize = kwargs.get('dsize', 224)33        self.timer = Timer()34 35        if onnx_provider.lower() == 'cuda':36            self.session = onnxruntime.InferenceSession(37                ckpt_path, providers=[38                    ('CUDAExecutionProvider', {'device_id': device_id})39                ]40            )41        else:42            opts = onnxruntime.SessionOptions()43            opts.intra_op_num_threads = 4  # 默认线程数为 444            self.session = onnxruntime.InferenceSession(45                ckpt_path, providers=['CPUExecutionProvider'],46                sess_options=opts47            )48 49    def _run(self, inp):50        out = self.session.run(None, {'input': inp})51        return out52 53    def run(self, img_rgb: np.ndarray, lmk=None):54        if lmk is not None:55            crop_dct = crop_image(img_rgb, lmk, dsize=self.dsize, scale=1.5, vy_ratio=-0.1)56            img_crop_rgb = crop_dct['img_crop']57        else:58            img_crop_rgb = cv2.resize(img_rgb, (self.dsize, self.dsize))59            scale = max(img_rgb.shape[:2]) / self.dsize60            crop_dct = {61                'M_c2o': np.array([62                    [scale, 0., 0.],63                    [0., scale, 0.],64                    [0., 0., 1.],65                ], dtype=np.float32),66            }67 68        inp = (img_crop_rgb.astype(np.float32) / 255.).transpose(2, 0, 1)[None, ...]  # HxWx3 (BGR) -> 1x3xHxW (RGB!)69 70        out_lst = self._run(inp)71        out_pts = out_lst[2]72 73        pts = to_ndarray(out_pts[0]).reshape(-1, 2) * self.dsize  # scale to 0-22474        pts = _transform_pts(pts, M=crop_dct['M_c2o'])75 76        return {77            'pts': pts,  # 2d landmarks 203 points78        }79 80    def warmup(self):81        # 构造dummy image进行warmup82        self.timer.tic()83 84        dummy_image = np.zeros((1, 3, self.dsize, self.dsize), dtype=np.float32)85 86        _ = self._run(dummy_image)87 88        elapse = self.timer.toc()89        rlog(f'LandmarkRunner warmup time: {elapse:.3f}s')90