moralec/MagicQuill

1import webcolors2import random3from collections import Counter4import numpy as np5from torchvision import transforms6import cv2  # OpenCV7import torch8import warnings9import os10 11 12 13def HWC3(x):14    assert x.dtype == np.uint815    if x.ndim == 2:16        x = x[:, :, None]17    assert x.ndim == 318    H, W, C = x.shape19    assert C == 1 or C == 3 or C == 420    if C == 3:21        return x22    if C == 1:23        return np.concatenate([x, x, x], axis=2)24    if C == 4:25        color = x[:, :, 0:3].astype(np.float32)26        alpha = x[:, :, 3:4].astype(np.float32) / 255.027        y = color * alpha + 255.0 * (1.0 - alpha)28        y = y.clip(0, 255).astype(np.uint8)29        return y30    31def common_input_validate(input_image, output_type, **kwargs):32    if "img" in kwargs:33            warnings.warn("img is deprecated, please use `input_image=...` instead.", DeprecationWarning)34            input_image = kwargs.pop("img")35    36    if "return_pil" in kwargs:37            warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning)38            output_type = "pil" if kwargs["return_pil"] else "np"39    40    if type(output_type) is bool:41        warnings.warn("Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions")42        if output_type:43            output_type = "pil"44 45    if input_image is None:46        raise ValueError("input_image must be defined.")47 48    if not isinstance(input_image, np.ndarray):49        input_image = np.array(input_image, dtype=np.uint8)50        output_type = output_type or "pil"51    else:52        output_type = output_type or "np"53    54    return (input_image, output_type)55 56def cv2_resize_shortest_edge(image, size):57    h, w = image.shape[:2]58    if h < w:59        new_h = size60        new_w = int(round(w / h * size))61    else:62        new_w = size63        new_h = int(round(h / w * size))64    resized_image = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_AREA)65    return resized_image66 67def apply_color(img, res=512):68    img = cv2_resize_shortest_edge(img, res)69    h, w = img.shape[:2]70 71    input_img_color = cv2.resize(img, (w//64, h//64), interpolation=cv2.INTER_CUBIC)  72    input_img_color = cv2.resize(input_img_color, (w, h), interpolation=cv2.INTER_NEAREST)73    return input_img_color74 75UPSCALE_METHODS = ["INTER_NEAREST", "INTER_LINEAR", "INTER_AREA", "INTER_CUBIC", "INTER_LANCZOS4"]76def get_upscale_method(method_str):77    assert method_str in UPSCALE_METHODS, f"Method {method_str} not found in {UPSCALE_METHODS}"78    return getattr(cv2, method_str)79 80def pad64(x):81    return int(np.ceil(float(x) / 64.0) * 64 - x)82 83def safer_memory(x):84    # Fix many MAC/AMD problems85    return np.ascontiguousarray(x.copy()).copy()86 87def resize_image_with_pad(input_image, resolution, upscale_method = "", skip_hwc3=False, mode='edge'):88    if skip_hwc3:89        img = input_image90    else:91        img = HWC3(input_image)92    H_raw, W_raw, _ = img.shape93    if resolution == 0:94        return img, lambda x: x95    k = float(resolution) / float(min(H_raw, W_raw))96    H_target = int(np.round(float(H_raw) * k))97    W_target = int(np.round(float(W_raw) * k))98    img = cv2.resize(img, (W_target, H_target), interpolation=get_upscale_method(upscale_method) if k > 1 else cv2.INTER_AREA)99    H_pad, W_pad = pad64(H_target), pad64(W_target)100    img_padded = np.pad(img, [[0, H_pad], [0, W_pad], [0, 0]], mode=mode)101 102    def remove_pad(x):103        return safer_memory(x[:H_target, :W_target, ...])104 105    return safer_memory(img_padded), remove_pad106 107def draw_contour(img, mask):108    mask_np = mask.numpy().astype(np.uint8) * 255109    img_np = img.numpy()110    img_np = img_np.astype(np.uint8)111    img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)112 113    kernel = np.ones((5, 5), np.uint8)114    mask_dilated = cv2.dilate(mask_np, kernel, iterations=3)115    contours, _ = cv2.findContours(mask_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)116    for contour in contours:117        cv2.drawContours(img_bgr, [contour], -1, (0, 0, 255), thickness=10)118    img_np = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)119    transform = transforms.ToTensor()120    img_tensor = transform(img_np)121 122    img_tensor = img_tensor.permute(1, 2, 0)123 124    return img_tensor.unsqueeze(0)125 126def get_colored_contour(img1, img2, threshold=10):127    diff = torch.abs(img1 - img2).float()128    diff_gray = torch.mean(diff, dim=-1)129    mask = diff_gray > threshold130 131    return draw_contour(img2, mask), mask132 133def closest_colour(requested_colour):134    min_colours = {}135    for key, name in webcolors.CSS3_HEX_TO_NAMES.items():136        r_c, g_c, b_c = webcolors.hex_to_rgb(key)137        rd = (r_c - requested_colour[0].item()) ** 2138        gd = (g_c - requested_colour[1].item()) ** 2139        bd = (b_c - requested_colour[2].item()) ** 2140        min_colours[(rd + gd + bd)] = name141    return min_colours[min(min_colours.keys())]142 143def rgb_to_name(rgb_tuple):144    try:145        return webcolors.rgb_to_name(rgb_tuple)146    except ValueError:147        closest_name = closest_colour(rgb_tuple)148        return closest_name149 150def find_different_colors(img1, img2, threshold=10):151    img1 = img1.to(torch.uint8)152    img2 = img2.to(torch.uint8)153    diff = torch.abs(img1 - img2).float().mean(dim=-1)154    diff_mask = diff > threshold155    diff_indices = torch.nonzero(diff_mask, as_tuple=True)156 157    if len(diff_indices[0]) > 100:158        sampled_indices = random.sample(range(len(diff_indices[0])), 100)159        sampled_diff_indices = (diff_indices[0][sampled_indices], diff_indices[1][sampled_indices])160    else:161        sampled_diff_indices = diff_indices162 163    diff_colors = img2[sampled_diff_indices[0], sampled_diff_indices[1], :]164    color_names = [rgb_to_name(tuple(color)) for color in diff_colors]165    name_counter = Counter(color_names)166    filtered_colors = {name: count for name, count in name_counter.items() if count > 10}167    sorted_color_names = [name for name, count in sorted(filtered_colors.items(), key=lambda item: item[1], reverse=True)]168    if len(sorted_color_names) >= 3:169        return "colorful"170    unique_color_names_str = ', '.join(sorted_color_names)171    return unique_color_names_str172 173def get_bounding_box_from_mask(mask, padded=False):174    # Ensure the mask is a binary mask (0s and 1s)175    mask = mask.squeeze()176    rows, cols = torch.where(mask > 0.5)177    if len(rows) == 0 or len(cols) == 0:178        return (0, 0, 0, 0)179    height, width = mask.shape180    if padded:181        padded_size = max(width, height)182        if width < height:183            offset_x = (padded_size - width) / 2184            offset_y = 0185        else:186            offset_y = (padded_size - height) / 2187            offset_x = 0188        # Find the bounding box coordinates189        top_left_x = round(float((torch.min(cols).item() + offset_x) / padded_size), 3)190        bottom_right_x = round(float((torch.max(cols).item() + offset_x) / padded_size), 3)191        top_left_y = round(float((torch.min(rows).item() + offset_y) / padded_size), 3)192        bottom_right_y = round(float((torch.max(rows).item() + offset_y) / padded_size), 3)193    else:194        offset_x = 0195        offset_y = 0196 197        top_left_x = round(float(torch.min(cols).item() / width), 3)198        bottom_right_x = round(float(torch.max(cols).item() / width), 3)199        top_left_y = round(float(torch.min(rows).item() / height), 3)200        bottom_right_y = round(float(torch.max(rows).item() / height), 3)201 202    203    return (top_left_x, top_left_y, bottom_right_x, bottom_right_y)