Shellbrady/LivePortrait112w

1# coding: utf-82 3import os4from glob import glob5import os.path as osp6import imageio7import numpy as np8import cv2; cv2.setNumThreads(0); cv2.ocl.setUseOpenCL(False)9 10 11def load_image_rgb(image_path: str):12    if not osp.exists(image_path):13        raise FileNotFoundError(f"Image not found: {image_path}")14    img = cv2.imread(image_path, cv2.IMREAD_COLOR)15    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)16 17 18def load_driving_info(driving_info):19    driving_video_ori = []20 21    def load_images_from_directory(directory):22        image_paths = sorted(glob(osp.join(directory, '*.png')) + glob(osp.join(directory, '*.jpg')))23        return [load_image_rgb(im_path) for im_path in image_paths]24 25    def load_images_from_video(file_path):26        reader = imageio.get_reader(file_path)27        return [image for idx, image in enumerate(reader)]28 29    if osp.isdir(driving_info):30        driving_video_ori = load_images_from_directory(driving_info)31    elif osp.isfile(driving_info):32        driving_video_ori = load_images_from_video(driving_info)33 34    return driving_video_ori35 36 37def contiguous(obj):38    if not obj.flags.c_contiguous:39        obj = obj.copy(order="C")40    return obj41 42 43def resize_to_limit(img: np.ndarray, max_dim=1920, n=2):44    """45    ajust the size of the image so that the maximum dimension does not exceed max_dim, and the width and the height of the image are multiples of n.46    :param img: the image to be processed.47    :param max_dim: the maximum dimension constraint.48    :param n: the number that needs to be multiples of.49    :return: the adjusted image.50    """51    h, w = img.shape[:2]52 53    # ajust the size of the image according to the maximum dimension54    if max_dim > 0 and max(h, w) > max_dim:55        if h > w:56            new_h = max_dim57            new_w = int(w * (max_dim / h))58        else:59            new_w = max_dim60            new_h = int(h * (max_dim / w))61        img = cv2.resize(img, (new_w, new_h))62 63    # ensure that the image dimensions are multiples of n64    n = max(n, 1)65    new_h = img.shape[0] - (img.shape[0] % n)66    new_w = img.shape[1] - (img.shape[1] % n)67 68    if new_h == 0 or new_w == 0:69        # when the width or height is less than n, no need to process70        return img71 72    if new_h != img.shape[0] or new_w != img.shape[1]:73        img = img[:new_h, :new_w]74 75    return img76 77 78def load_img_online(obj, mode="bgr", **kwargs):79    max_dim = kwargs.get("max_dim", 1920)80    n = kwargs.get("n", 2)81    if isinstance(obj, str):82        if mode.lower() == "gray":83            img = cv2.imread(obj, cv2.IMREAD_GRAYSCALE)84        else:85            img = cv2.imread(obj, cv2.IMREAD_COLOR)86    else:87        img = obj88 89    # Resize image to satisfy constraints90    img = resize_to_limit(img, max_dim=max_dim, n=n)91 92    if mode.lower() == "bgr":93        return contiguous(img)94    elif mode.lower() == "rgb":95        return contiguous(img[..., ::-1])96    else:97        raise Exception(f"Unknown mode {mode}")98