Spaces:
Runtime error
Runtime error
File size: 1,584 Bytes
c19ca42 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
import cv2
import numpy as np
from ..supported_preprocessor import Preprocessor, PreprocessorParameter
from ..utils import resize_image_with_pad
class PreprocessorLamaInpaint(Preprocessor):
def __init__(self):
super().__init__(name="inpaint_only+lama")
self.tags = ["Inpaint"]
self.returns_image = True
self.model = None
self.slider_resolution = PreprocessorParameter(visible=False)
def __call__(
self,
input_image,
resolution,
slider_1=None,
slider_2=None,
slider_3=None,
**kwargs
):
img = input_image
H, W, C = img.shape
raw_color = img[:, :, 0:3].copy()
raw_mask = img[:, :, 3:4].copy()
res = 256 # Always use 256 since lama is trained on 256
img_res, remove_pad = resize_image_with_pad(img, res)
if self.model is None:
from annotator.lama import LamaInpainting
self.model = LamaInpainting()
# applied auto inversion
prd_color = self.model(img_res)
self.model.unload_model()
prd_color = remove_pad(prd_color)
prd_color = cv2.resize(prd_color, (W, H))
alpha = raw_mask.astype(np.float32) / 255.0
fin_color = prd_color.astype(np.float32) * alpha + raw_color.astype(
np.float32
) * (1 - alpha)
fin_color = fin_color.clip(0, 255).astype(np.uint8)
result = np.concatenate([fin_color, raw_mask], axis=2)
return result
Preprocessor.add_supported_preprocessor(PreprocessorLamaInpaint())
|