_
File size: 744 Bytes
da3eeba
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
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):
    name = "cv2"
    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