svjack/Kolors-Controlnet_and_IPA

1import random2 3import numpy as np4import cv25import os6import PIL7 8annotator_ckpts_path = os.path.join(os.path.dirname(__file__), 'ckpts')9 10def HWC3(x):11    assert x.dtype == np.uint812    if x.ndim == 2:13        x = x[:, :, None]14    assert x.ndim == 315    H, W, C = x.shape16    assert C == 1 or C == 3 or C == 417    if C == 3:18        return x19    if C == 1:20        return np.concatenate([x, x, x], axis=2)21    if C == 4:22        color = x[:, :, 0:3].astype(np.float32)23        alpha = x[:, :, 3:4].astype(np.float32) / 255.024        y = color * alpha + 255.0 * (1.0 - alpha)25        y = y.clip(0, 255).astype(np.uint8)26        return y27 28 29def resize_image(input_image, resolution, short = False, interpolation=None):30    if isinstance(input_image,PIL.Image.Image):31        mode = 'pil'32        W,H = input_image.size33 34    elif isinstance(input_image,np.ndarray):35        mode = 'cv2'36        H, W, _ = input_image.shape37 38    H = float(H)39    W = float(W)40    if short:41        k = float(resolution) / min(H, W) # k>1 放大， k<1 缩小42    else:43        k = float(resolution) / max(H, W) # k>1 放大， k<1 缩小44    H *= k 45    W *= k46    H = int(np.round(H / 64.0)) * 6447    W = int(np.round(W / 64.0)) * 6448    49    if mode == 'cv2':50        if interpolation is None:51            interpolation = cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA52        img = cv2.resize(input_image, (W, H), interpolation=interpolation)53 54    elif mode == 'pil':55        if interpolation is None:56            interpolation = PIL.Image.LANCZOS if k > 1 else PIL.Image.BILINEAR57        img = input_image.resize((W, H), resample=interpolation)58    59    return img60 61# def resize_image(input_image, resolution):62#     H, W, C = input_image.shape63#     H = float(H)64#     W = float(W)65#     k = float(resolution) / min(H, W)66#     H *= k67#     W *= k68#     H = int(np.round(H / 64.0)) * 6469#     W = int(np.round(W / 64.0)) * 6470#     img = cv2.resize(input_image, (W, H), interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA)71#     return img72 73 74def nms(x, t, s):75    x = cv2.GaussianBlur(x.astype(np.float32), (0, 0), s)76 77    f1 = np.array([[0, 0, 0], [1, 1, 1], [0, 0, 0]], dtype=np.uint8)78    f2 = np.array([[0, 1, 0], [0, 1, 0], [0, 1, 0]], dtype=np.uint8)79    f3 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.uint8)80    f4 = np.array([[0, 0, 1], [0, 1, 0], [1, 0, 0]], dtype=np.uint8)81 82    y = np.zeros_like(x)83 84    for f in [f1, f2, f3, f4]:85        np.putmask(y, cv2.dilate(x, kernel=f) == x, x)86 87    z = np.zeros_like(y, dtype=np.uint8)88    z[y > t] = 25589    return z90 91 92def make_noise_disk(H, W, C, F):93    noise = np.random.uniform(low=0, high=1, size=((H // F) + 2, (W // F) + 2, C))94    noise = cv2.resize(noise, (W + 2 * F, H + 2 * F), interpolation=cv2.INTER_CUBIC)95    noise = noise[F: F + H, F: F + W]96    noise -= np.min(noise)97    noise /= np.max(noise)98    if C == 1:99        noise = noise[:, :, None]100    return noise101 102 103def min_max_norm(x):104    x -= np.min(x)105    x /= np.maximum(np.max(x), 1e-5)106    return x107 108 109def safe_step(x, step=2):110    y = x.astype(np.float32) * float(step + 1)111    y = y.astype(np.int32).astype(np.float32) / float(step)112    return y113 114 115def img2mask(img, H, W, low=10, high=90):116    assert img.ndim == 3 or img.ndim == 2117    assert img.dtype == np.uint8118 119    if img.ndim == 3:120        y = img[:, :, random.randrange(0, img.shape[2])]121    else:122        y = img123 124    y = cv2.resize(y, (W, H), interpolation=cv2.INTER_CUBIC)125 126    if random.uniform(0, 1) < 0.5:127        y = 255 - y128 129    return y < np.percentile(y, random.randrange(low, high))130