Add gaussian blur to mask before blending original image with result

This commit is contained in:
Anders Haglund 2022-11-10 09:45:20 -08:00
parent 9144e101a1
commit 543473bac5

View File

@ -3,6 +3,7 @@ from typing import Optional
import cv2 import cv2
import torch import torch
import numpy as np
from loguru import logger from loguru import logger
from lama_cleaner.helper import boxes_from_mask, resize_max_size, pad_img_to_modulo from lama_cleaner.helper import boxes_from_mask, resize_max_size, pad_img_to_modulo
@ -55,8 +56,12 @@ class InpaintModel:
result = self.forward(pad_image, pad_mask, config) result = self.forward(pad_image, pad_mask, config)
result = result[0:origin_height, 0:origin_width, :] result = result[0:origin_height, 0:origin_width, :]
original_pixel_indices = mask < 127 if config.sd_mask_blur != 0:
result[original_pixel_indices] = image[:, :, ::-1][original_pixel_indices] k = 2 * config.sd_mask_blur + 1
mask = cv2.GaussianBlur(mask, (k, k), 0)
mask = mask[:, :, np.newaxis]
result = result * (mask / 255) + image[:, :, ::-1] * (1 - (mask / 255))
return result return result
@torch.no_grad() @torch.no_grad()