Add OpenCV2 model
This commit is contained in:
19
lama_cleaner/model/opencv2.py
Normal file
19
lama_cleaner/model/opencv2.py
Normal 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
|
||||
Reference in New Issue
Block a user