georgefen/Face-Landmark-ControlNet

1"""Utils for monoDepth."""2import sys3import re4import numpy as np5import cv26import torch7 8 9def read_pfm(path):10    """Read pfm file.11 12    Args:13        path (str): path to file14 15    Returns:16        tuple: (data, scale)17    """18    with open(path, "rb") as file:19 20        color = None21        width = None22        height = None23        scale = None24        endian = None25 26        header = file.readline().rstrip()27        if header.decode("ascii") == "PF":28            color = True29        elif header.decode("ascii") == "Pf":30            color = False31        else:32            raise Exception("Not a PFM file: " + path)33 34        dim_match = re.match(r"^(\d+)\s(\d+)\s$", file.readline().decode("ascii"))35        if dim_match:36            width, height = list(map(int, dim_match.groups()))37        else:38            raise Exception("Malformed PFM header.")39 40        scale = float(file.readline().decode("ascii").rstrip())41        if scale < 0:42            # little-endian43            endian = "<"44            scale = -scale45        else:46            # big-endian47            endian = ">"48 49        data = np.fromfile(file, endian + "f")50        shape = (height, width, 3) if color else (height, width)51 52        data = np.reshape(data, shape)53        data = np.flipud(data)54 55        return data, scale56 57 58def write_pfm(path, image, scale=1):59    """Write pfm file.60 61    Args:62        path (str): pathto file63        image (array): data64        scale (int, optional): Scale. Defaults to 1.65    """66 67    with open(path, "wb") as file:68        color = None69 70        if image.dtype.name != "float32":71            raise Exception("Image dtype must be float32.")72 73        image = np.flipud(image)74 75        if len(image.shape) == 3 and image.shape[2] == 3:  # color image76            color = True77        elif (78            len(image.shape) == 2 or len(image.shape) == 3 and image.shape[2] == 179        ):  # greyscale80            color = False81        else:82            raise Exception("Image must have H x W x 3, H x W x 1 or H x W dimensions.")83 84        file.write("PF\n" if color else "Pf\n".encode())85        file.write("%d %d\n".encode() % (image.shape[1], image.shape[0]))86 87        endian = image.dtype.byteorder88 89        if endian == "<" or endian == "=" and sys.byteorder == "little":90            scale = -scale91 92        file.write("%f\n".encode() % scale)93 94        image.tofile(file)95 96 97def read_image(path):98    """Read image and output RGB image (0-1).99 100    Args:101        path (str): path to file102 103    Returns:104        array: RGB image (0-1)105    """106    img = cv2.imread(path)107 108    if img.ndim == 2:109        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)110 111    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) / 255.0112 113    return img114 115 116def resize_image(img):117    """Resize image and make it fit for network.118 119    Args:120        img (array): image121 122    Returns:123        tensor: data ready for network124    """125    height_orig = img.shape[0]126    width_orig = img.shape[1]127 128    if width_orig > height_orig:129        scale = width_orig / 384130    else:131        scale = height_orig / 384132 133    height = (np.ceil(height_orig / scale / 32) * 32).astype(int)134    width = (np.ceil(width_orig / scale / 32) * 32).astype(int)135 136    img_resized = cv2.resize(img, (width, height), interpolation=cv2.INTER_AREA)137 138    img_resized = (139        torch.from_numpy(np.transpose(img_resized, (2, 0, 1))).contiguous().float()140    )141    img_resized = img_resized.unsqueeze(0)142 143    return img_resized144 145 146def resize_depth(depth, width, height):147    """Resize depth map and bring to CPU (numpy).148 149    Args:150        depth (tensor): depth151        width (int): image width152        height (int): image height153 154    Returns:155        array: processed depth156    """157    depth = torch.squeeze(depth[0, :, :, :]).to("cpu")158 159    depth_resized = cv2.resize(160        depth.numpy(), (width, height), interpolation=cv2.INTER_CUBIC161    )162 163    return depth_resized164 165def write_depth(path, depth, bits=1):166    """Write depth map to pfm and png file.167 168    Args:169        path (str): filepath without extension170        depth (array): depth171    """172    write_pfm(path + ".pfm", depth.astype(np.float32))173 174    depth_min = depth.min()175    depth_max = depth.max()176 177    max_val = (2**(8*bits))-1178 179    if depth_max - depth_min > np.finfo("float").eps:180        out = max_val * (depth - depth_min) / (depth_max - depth_min)181    else:182        out = np.zeros(depth.shape, dtype=depth.type)183 184    if bits == 1:185        cv2.imwrite(path + ".png", out.astype("uint8"))186    elif bits == 2:187        cv2.imwrite(path + ".png", out.astype("uint16"))188 189    return190