30 lines
735 B
Python
30 lines
735 B
Python
import cv2
|
|
from .base import InpaintModel
|
|
from iopaint.schema import InpaintRequest
|
|
|
|
flag_map = {"INPAINT_NS": cv2.INPAINT_NS, "INPAINT_TELEA": cv2.INPAINT_TELEA}
|
|
|
|
|
|
class OpenCV2(InpaintModel):
|
|
name = "cv2"
|
|
pad_mod = 1
|
|
is_erase_model = True
|
|
|
|
@staticmethod
|
|
def is_downloaded() -> bool:
|
|
return True
|
|
|
|
def forward(self, image, mask, config: InpaintRequest):
|
|
"""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
|