File size: 648 Bytes
ed01507 |
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 |
import cv2
from lama_cleaner.model.base import InpaintModel
from lama_cleaner.schema import Config
flag_map = {
"INPAINT_NS": cv2.INPAINT_NS,
"INPAINT_TELEA": cv2.INPAINT_TELEA
}
class OpenCV2(InpaintModel):
pad_mod = 1
@staticmethod
def is_downloaded() -> bool:
return True
def forward(self, image, mask, config: Config):
"""Input image and output image have same size
image: [H, W, C] RGB
mask: [H, W, 1]
return: BGR IMAGE
"""
cur_res = cv2.inpaint(image[:,:,::-1], mask, inpaintRadius=config.cv2_radius, flags=flag_map[config.cv2_flag])
return cur_res
|