|
|
|
|
|
import os
|
|
from glob import glob
|
|
import os.path as osp
|
|
import imageio
|
|
import numpy as np
|
|
import pickle
|
|
import cv2; cv2.setNumThreads(0); cv2.ocl.setUseOpenCL(False)
|
|
|
|
from .helper import mkdir, suffix
|
|
|
|
|
|
def load_image_rgb(image_path: str):
|
|
if not osp.exists(image_path):
|
|
raise FileNotFoundError(f"Image not found: {image_path}")
|
|
img = cv2.imread(image_path, cv2.IMREAD_COLOR)
|
|
return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
|
|
|
|
|
def load_driving_info(driving_info):
|
|
driving_video_ori = []
|
|
|
|
def load_images_from_directory(directory):
|
|
image_paths = sorted(glob(osp.join(directory, '*.png')) + glob(osp.join(directory, '*.jpg')))
|
|
return [load_image_rgb(im_path) for im_path in image_paths]
|
|
|
|
def load_images_from_video(file_path):
|
|
reader = imageio.get_reader(file_path, "ffmpeg")
|
|
return [image for _, image in enumerate(reader)]
|
|
|
|
if osp.isdir(driving_info):
|
|
driving_video_ori = load_images_from_directory(driving_info)
|
|
elif osp.isfile(driving_info):
|
|
driving_video_ori = load_images_from_video(driving_info)
|
|
|
|
return driving_video_ori
|
|
|
|
|
|
def contiguous(obj):
|
|
if not obj.flags.c_contiguous:
|
|
obj = obj.copy(order="C")
|
|
return obj
|
|
|
|
|
|
def resize_to_limit(img: np.ndarray, max_dim=1920, division=2):
|
|
"""
|
|
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.
|
|
:param img: the image to be processed.
|
|
:param max_dim: the maximum dimension constraint.
|
|
:param n: the number that needs to be multiples of.
|
|
:return: the adjusted image.
|
|
"""
|
|
h, w = img.shape[:2]
|
|
|
|
|
|
if max_dim > 0 and max(h, w) > max_dim:
|
|
if h > w:
|
|
new_h = max_dim
|
|
new_w = int(w * (max_dim / h))
|
|
else:
|
|
new_w = max_dim
|
|
new_h = int(h * (max_dim / w))
|
|
img = cv2.resize(img, (new_w, new_h))
|
|
|
|
|
|
division = max(division, 1)
|
|
new_h = img.shape[0] - (img.shape[0] % division)
|
|
new_w = img.shape[1] - (img.shape[1] % division)
|
|
|
|
if new_h == 0 or new_w == 0:
|
|
|
|
return img
|
|
|
|
if new_h != img.shape[0] or new_w != img.shape[1]:
|
|
img = img[:new_h, :new_w]
|
|
|
|
return img
|
|
|
|
|
|
def load_img_online(obj, mode="bgr", **kwargs):
|
|
max_dim = kwargs.get("max_dim", 1920)
|
|
n = kwargs.get("n", 2)
|
|
if isinstance(obj, str):
|
|
if mode.lower() == "gray":
|
|
img = cv2.imread(obj, cv2.IMREAD_GRAYSCALE)
|
|
else:
|
|
img = cv2.imread(obj, cv2.IMREAD_COLOR)
|
|
else:
|
|
img = obj
|
|
|
|
|
|
img = resize_to_limit(img, max_dim=max_dim, division=n)
|
|
|
|
if mode.lower() == "bgr":
|
|
return contiguous(img)
|
|
elif mode.lower() == "rgb":
|
|
return contiguous(img[..., ::-1])
|
|
else:
|
|
raise Exception(f"Unknown mode {mode}")
|
|
|
|
|
|
def load(fp):
|
|
suffix_ = suffix(fp)
|
|
|
|
if suffix_ == "npy":
|
|
return np.load(fp)
|
|
elif suffix_ == "pkl":
|
|
return pickle.load(open(fp, "rb"))
|
|
else:
|
|
raise Exception(f"Unknown type: {suffix}")
|
|
|
|
|
|
def dump(wfp, obj):
|
|
wd = osp.split(wfp)[0]
|
|
if wd != "" and not osp.exists(wd):
|
|
mkdir(wd)
|
|
|
|
_suffix = suffix(wfp)
|
|
if _suffix == "npy":
|
|
np.save(wfp, obj)
|
|
elif _suffix == "pkl":
|
|
pickle.dump(obj, open(wfp, "wb"))
|
|
else:
|
|
raise Exception("Unknown type: {}".format(_suffix))
|
|
|
