coreml-community/ControlNet-v1-1-Annotators-cpu

1import random2 3import numpy as np4import cv25import os6 7 8annotator_ckpts_path = os.path.join(os.path.dirname(__file__), 'ckpts')9 10 11def HWC3(x):12    assert x.dtype == np.uint813    if x.ndim == 2:14        x = x[:, :, None]15    assert x.ndim == 316    H, W, C = x.shape17    assert C == 1 or C == 3 or C == 418    if C == 3:19        return x20    if C == 1:21        return np.concatenate([x, x, x], axis=2)22    if C == 4:23        color = x[:, :, 0:3].astype(np.float32)24        alpha = x[:, :, 3:4].astype(np.float32) / 255.025        y = color * alpha + 255.0 * (1.0 - alpha)26        y = y.clip(0, 255).astype(np.uint8)27        return y28 29 30def resize_image(input_image, resolution):31    H, W, C = input_image.shape32    H = float(H)33    W = float(W)34    k = float(resolution) / min(H, W)35    H *= k36    W *= k37    H = int(np.round(H / 64.0)) * 6438    W = int(np.round(W / 64.0)) * 6439    img = cv2.resize(input_image, (W, H), interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA)40    return img41 42 43def nms(x, t, s):44    x = cv2.GaussianBlur(x.astype(np.float32), (0, 0), s)45 46    f1 = np.array([[0, 0, 0], [1, 1, 1], [0, 0, 0]], dtype=np.uint8)47    f2 = np.array([[0, 1, 0], [0, 1, 0], [0, 1, 0]], dtype=np.uint8)48    f3 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.uint8)49    f4 = np.array([[0, 0, 1], [0, 1, 0], [1, 0, 0]], dtype=np.uint8)50 51    y = np.zeros_like(x)52 53    for f in [f1, f2, f3, f4]:54        np.putmask(y, cv2.dilate(x, kernel=f) == x, x)55 56    z = np.zeros_like(y, dtype=np.uint8)57    z[y > t] = 25558    return z59 60 61def make_noise_disk(H, W, C, F):62    noise = np.random.uniform(low=0, high=1, size=((H // F) + 2, (W // F) + 2, C))63    noise = cv2.resize(noise, (W + 2 * F, H + 2 * F), interpolation=cv2.INTER_CUBIC)64    noise = noise[F: F + H, F: F + W]65    noise -= np.min(noise)66    noise /= np.max(noise)67    if C == 1:68        noise = noise[:, :, None]69    return noise70 71 72def min_max_norm(x):73    x -= np.min(x)74    x /= np.maximum(np.max(x), 1e-5)75    return x76 77 78def safe_step(x, step=2):79    y = x.astype(np.float32) * float(step + 1)80    y = y.astype(np.int32).astype(np.float32) / float(step)81    return y82 83 84def img2mask(img, H, W, low=10, high=90):85    assert img.ndim == 3 or img.ndim == 286    assert img.dtype == np.uint887 88    if img.ndim == 3:89        y = img[:, :, random.randrange(0, img.shape[2])]90    else:91        y = img92 93    y = cv2.resize(y, (W, H), interpolation=cv2.INTER_CUBIC)94 95    if random.uniform(0, 1) < 0.5:96        y = 255 - y97 98    return y < np.percentile(y, random.randrange(low, high))99