| import os | |
| import random | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| annotator_ckpts_path = os.path.join(os.path.dirname(__file__), 'ckpts') | |
| def HWC3(x): | |
| assert x.dtype == np.uint8 | |
| if x.ndim == 2: | |
| x = x[:, :, None] | |
| assert x.ndim == 3 | |
| H, W, C = x.shape | |
| assert C == 1 or C == 3 or C == 4 | |
| if C == 3: | |
| return x | |
| if C == 1: | |
| return np.concatenate([x, x, x], axis=2) | |
| if C == 4: | |
| color = x[:, :, 0:3].astype(np.float32) | |
| alpha = x[:, :, 3:4].astype(np.float32) / 255.0 | |
| y = color * alpha + 255.0 * (1.0 - alpha) | |
| y = y.clip(0, 255).astype(np.uint8) | |
| return y | |
| def make_noise_disk(H, W, C, F): | |
| noise = np.random.uniform(low=0, high=1, size=((H // F) + 2, (W // F) + 2, C)) | |
| noise = cv2.resize(noise, (W + 2 * F, H + 2 * F), interpolation=cv2.INTER_CUBIC) | |
| noise = noise[F: F + H, F: F + W] | |
| noise -= np.min(noise) | |
| noise /= np.max(noise) | |
| if C == 1: | |
| noise = noise[:, :, None] | |
| return noise | |
| def nms(x, t, s): | |
| x = cv2.GaussianBlur(x.astype(np.float32), (0, 0), s) | |
| f1 = np.array([[0, 0, 0], [1, 1, 1], [0, 0, 0]], dtype=np.uint8) | |
| f2 = np.array([[0, 1, 0], [0, 1, 0], [0, 1, 0]], dtype=np.uint8) | |
| f3 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.uint8) | |
| f4 = np.array([[0, 0, 1], [0, 1, 0], [1, 0, 0]], dtype=np.uint8) | |
| y = np.zeros_like(x) | |
| for f in [f1, f2, f3, f4]: | |
| np.putmask(y, cv2.dilate(x, kernel=f) == x, x) | |
| z = np.zeros_like(y, dtype=np.uint8) | |
| z[y > t] = 255 | |
| return z | |
| def min_max_norm(x): | |
| x -= np.min(x) | |
| x /= np.maximum(np.max(x), 1e-5) | |
| return x | |
| def safe_step(x, step=2): | |
| y = x.astype(np.float32) * float(step + 1) | |
| y = y.astype(np.int32).astype(np.float32) / float(step) | |
| return y | |
| def img2mask(img, H, W, low=10, high=90): | |
| assert img.ndim == 3 or img.ndim == 2 | |
| assert img.dtype == np.uint8 | |
| if img.ndim == 3: | |
| y = img[:, :, random.randrange(0, img.shape[2])] | |
| else: | |
| y = img | |
| y = cv2.resize(y, (W, H), interpolation=cv2.INTER_CUBIC) | |
| if random.uniform(0, 1) < 0.5: | |
| y = 255 - y | |
| return y < np.percentile(y, random.randrange(low, high)) | |
| def resize_image(input_image, resolution): | |
| H, W, C = input_image.shape | |
| H = float(H) | |
| W = float(W) | |
| k = float(resolution) / min(H, W) | |
| H *= k | |
| W *= k | |
| H = int(np.round(H / 64.0)) * 64 | |
| W = int(np.round(W / 64.0)) * 64 | |
| img = cv2.resize(input_image, (W, H), interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA) | |
| return img | |
| def torch_gc(): | |
| if torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |
| torch.cuda.ipc_collect() | |
| def ade_palette(): | |
| """ADE20K palette that maps each class to RGB values.""" | |
| return [[120, 120, 120], [180, 120, 120], [6, 230, 230], [80, 50, 50], | |
| [4, 200, 3], [120, 120, 80], [140, 140, 140], [204, 5, 255], | |
| [230, 230, 230], [4, 250, 7], [224, 5, 255], [235, 255, 7], | |
| [150, 5, 61], [120, 120, 70], [8, 255, 51], [255, 6, 82], | |
| [143, 255, 140], [204, 255, 4], [255, 51, 7], [204, 70, 3], | |
| [0, 102, 200], [61, 230, 250], [255, 6, 51], [11, 102, 255], | |
| [255, 7, 71], [255, 9, 224], [9, 7, 230], [220, 220, 220], | |
| [255, 9, 92], [112, 9, 255], [8, 255, 214], [7, 255, 224], | |
| [255, 184, 6], [10, 255, 71], [255, 41, 10], [7, 255, 255], | |
| [224, 255, 8], [102, 8, 255], [255, 61, 6], [255, 194, 7], | |
| [255, 122, 8], [0, 255, 20], [255, 8, 41], [255, 5, 153], | |
| [6, 51, 255], [235, 12, 255], [160, 150, 20], [0, 163, 255], | |
| [140, 140, 140], [250, 10, 15], [20, 255, 0], [31, 255, 0], | |
| [255, 31, 0], [255, 224, 0], [153, 255, 0], [0, 0, 255], | |
| [255, 71, 0], [0, 235, 255], [0, 173, 255], [31, 0, 255], | |
| [11, 200, 200], [255, 82, 0], [0, 255, 245], [0, 61, 255], | |
| [0, 255, 112], [0, 255, 133], [255, 0, 0], [255, 163, 0], | |
| [255, 102, 0], [194, 255, 0], [0, 143, 255], [51, 255, 0], | |
| [0, 82, 255], [0, 255, 41], [0, 255, 173], [10, 0, 255], | |
| [173, 255, 0], [0, 255, 153], [255, 92, 0], [255, 0, 255], | |
| [255, 0, 245], [255, 0, 102], [255, 173, 0], [255, 0, 20], | |
| [255, 184, 184], [0, 31, 255], [0, 255, 61], [0, 71, 255], | |
| [255, 0, 204], [0, 255, 194], [0, 255, 82], [0, 10, 255], | |
| [0, 112, 255], [51, 0, 255], [0, 194, 255], [0, 122, 255], | |
| [0, 255, 163], [255, 153, 0], [0, 255, 10], [255, 112, 0], | |
| [143, 255, 0], [82, 0, 255], [163, 255, 0], [255, 235, 0], | |
| [8, 184, 170], [133, 0, 255], [0, 255, 92], [184, 0, 255], | |
| [255, 0, 31], [0, 184, 255], [0, 214, 255], [255, 0, 112], | |
| [92, 255, 0], [0, 224, 255], [112, 224, 255], [70, 184, 160], | |
| [163, 0, 255], [153, 0, 255], [71, 255, 0], [255, 0, 163], | |
| [255, 204, 0], [255, 0, 143], [0, 255, 235], [133, 255, 0], | |
| [255, 0, 235], [245, 0, 255], [255, 0, 122], [255, 245, 0], | |
| [10, 190, 212], [214, 255, 0], [0, 204, 255], [20, 0, 255], | |
| [255, 255, 0], [0, 153, 255], [0, 41, 255], [0, 255, 204], | |
| [41, 0, 255], [41, 255, 0], [173, 0, 255], [0, 245, 255], | |
| [71, 0, 255], [122, 0, 255], [0, 255, 184], [0, 92, 255], | |
| [184, 255, 0], [0, 133, 255], [255, 214, 0], [25, 194, 194], | |
| [102, 255, 0], [92, 0, 255]] | |