MykolaL/StableDesign

1import gc2 3import numpy as np4from PIL import Image5import torch6from scipy.signal import fftconvolve7 8from palette import COLOR_MAPPING, COLOR_MAPPING_9 10 11def to_rgb(color: str) -> tuple:12    """Convert hex color to rgb.13    Args:14        color (str): hex color15    Returns:16        tuple: rgb color17    """18    return tuple(int(color[i:i+2], 16) for i in (1, 3, 5))19 20 21def map_colors(color: str) -> str:22    """Map color to hex value.23    Args:24        color (str): color name25    Returns:26        str: hex value27    """28    return COLOR_MAPPING[color]29 30 31def map_colors_rgb(color: tuple) -> str:32    return COLOR_MAPPING_RGB[color]33 34 35def convolution(mask: Image.Image, size=9) -> Image:36    """Method to blur the mask37    Args:38        mask (Image): masking image39        size (int, optional): size of the blur. Defaults to 9.40    Returns:41        Image: blurred mask42    """43    mask = np.array(mask.convert("L"))44    conv = np.ones((size, size)) / size**245    mask_blended = fftconvolve(mask, conv, 'same')46    mask_blended = mask_blended.astype(np.uint8).copy()47 48    border = size49 50    # replace borders with original values51    mask_blended[:border, :] = mask[:border, :]52    mask_blended[-border:, :] = mask[-border:, :]53    mask_blended[:, :border] = mask[:, :border]54    mask_blended[:, -border:] = mask[:, -border:]55 56    return Image.fromarray(mask_blended).convert("L")57 58 59def flush():60    gc.collect()61    torch.cuda.empty_cache()62 63 64def postprocess_image_masking(inpainted: Image, image: Image,65                              mask: Image) -> Image:66    """Method to postprocess the inpainted image67    Args:68        inpainted (Image): inpainted image69        image (Image): original image70        mask (Image): mask71    Returns:72        Image: inpainted image73    """74    final_inpainted = Image.composite(inpainted.convert("RGBA"),75                                      image.convert("RGBA"), mask)76    return final_inpainted.convert("RGB")77 78 79COLOR_NAMES = list(COLOR_MAPPING.keys())80COLOR_RGB = [to_rgb(k) for k in COLOR_MAPPING_.keys()] + [(0, 0, 0),81                                                          (255, 255, 255)]82INVERSE_COLORS = {v: to_rgb(k) for k, v in COLOR_MAPPING_.items()}83COLOR_MAPPING_RGB = {to_rgb(k): v for k, v in COLOR_MAPPING_.items()}84