Add OpenCV2 model

This commit is contained in:
Qing
2022-09-25 21:27:12 +08:00
parent bc98ea256a
commit 35cbbd653c
5 changed files with 88 additions and 3 deletions

View File

@@ -0,0 +1,19 @@
import cv2
from lama_cleaner.model.base import InpaintModel
from lama_cleaner.schema import Config
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=3, flags=cv2.INPAINT_TELEA)
return cur_res