|
import gc |
|
import torch |
|
from scipy.signal import fftconvolve |
|
from PIL import Image |
|
import numpy as np |
|
|
|
def flush(): |
|
gc.collect() |
|
torch.cuda.empty_cache() |
|
|
|
|
|
|
|
def convolution(mask: Image.Image, size=9) -> Image: |
|
"""Method to blur the mask |
|
Args: |
|
mask (Image): masking image |
|
size (int, optional): size of the blur. Defaults to 9. |
|
Returns: |
|
Image: blurred mask |
|
""" |
|
mask = np.array(mask.convert("L")) |
|
conv = np.ones((size, size)) / size**2 |
|
mask_blended = fftconvolve(mask, conv, 'same') |
|
mask_blended = mask_blended.astype(np.uint8).copy() |
|
|
|
border = size |
|
|
|
|
|
mask_blended[:border, :] = mask[:border, :] |
|
mask_blended[-border:, :] = mask[-border:, :] |
|
mask_blended[:, :border] = mask[:, :border] |
|
mask_blended[:, -border:] = mask[:, -border:] |
|
|
|
return Image.fromarray(mask_blended).convert("L") |
|
|
|
|
|
def postprocess_image_masking(inpainted: Image, image: Image, mask: Image) -> Image: |
|
"""Method to postprocess the inpainted image |
|
Args: |
|
inpainted (Image): inpainted image |
|
image (Image): original image |
|
mask (Image): mask |
|
Returns: |
|
Image: inpainted image |
|
""" |
|
final_inpainted = Image.composite(inpainted.convert("RGBA"), image.convert("RGBA"), mask) |
|
return final_inpainted.convert("RGB") |
|
|
