AideepImage/interior-design

1import gc2import torch3from scipy.signal import fftconvolve4from PIL import Image5import numpy as np6 7def flush():8    gc.collect()9    torch.cuda.empty_cache()10    11 12 13def convolution(mask: Image.Image, size=9) -> Image:14    """Method to blur the mask15    Args:16        mask (Image): masking image17        size (int, optional): size of the blur. Defaults to 9.18    Returns:19        Image: blurred mask20    """21    mask = np.array(mask.convert("L"))22    conv = np.ones((size, size)) / size**223    mask_blended = fftconvolve(mask, conv, 'same')24    mask_blended = mask_blended.astype(np.uint8).copy()25 26    border = size27 28    # replace borders with original values29    mask_blended[:border, :] = mask[:border, :]30    mask_blended[-border:, :] = mask[-border:, :]31    mask_blended[:, :border] = mask[:, :border]32    mask_blended[:, -border:] = mask[:, -border:]33 34    return Image.fromarray(mask_blended).convert("L")35 36 37def postprocess_image_masking(inpainted: Image, image: Image, mask: Image) -> Image:38    """Method to postprocess the inpainted image39    Args:40        inpainted (Image): inpainted image41        image (Image): original image42        mask (Image): mask43    Returns:44        Image: inpainted image45    """46    final_inpainted = Image.composite(inpainted.convert("RGBA"), image.convert("RGBA"), mask)47    return final_inpainted.convert("RGB")48