Add OpenCV2 model

lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx

@@ -1,6 +1,6 @@
 import React, { ReactNode } from 'react'
 import { useRecoilState } from 'recoil'
-import { AIModel, SDSampler, settingState } from '../../store/Atoms'
+import { AIModel, CV2Flag, SDSampler, settingState } from '../../store/Atoms'
 import Selector from '../shared/Selector'
 import { Switch, SwitchThumb } from '../shared/Switch'
 import Tooltip from '../shared/Tooltip'
@@ -133,6 +133,42 @@ function ModelSettingBlock() {
     )
   }
 
+  const renderOpenCV2Desc = () => {
+    return (
+      <>
+        <NumberInputSetting
+          title="Radius"
+          value={`${setting.cv2Radius}`}
+          desc="Radius of a circular neighborhood of each point inpainted that is considered by the algorithm."
+          onValue={value => {
+            const val = value.length === 0 ? 0 : parseInt(value, 10)
+            setSettingState(old => {
+              return { ...old, cv2Radius: val }
+            })
+          }}
+        />
+
+        <SettingBlock
+          className="sub-setting-block"
+          title="Flag"
+          desc="Inpainting method"
+          input={
+            <Selector
+              width={140}
+              value={setting.cv2Flag as string}
+              options={Object.values(CV2Flag)}
+              onChange={val => {
+                setSettingState(old => {
+                  return { ...old, cv2Flag: val as CV2Flag }
+                })
+              }}
+            />
+          }
+        />
+      </>
+    )
+  }
+
   const renderOptionDesc = (): ReactNode => {
     switch (setting.model) {
       case AIModel.LAMA:
@@ -147,6 +183,8 @@ function ModelSettingBlock() {
         return renderFCFModelDesc()
       case AIModel.SD14:
         return undefined
+      case AIModel.CV2:
+        return renderOpenCV2Desc()
       default:
         return <></>
     }
@@ -190,6 +228,12 @@ function ModelSettingBlock() {
           'https://ommer-lab.com/research/latent-diffusion-models/',
           'https://github.com/CompVis/stable-diffusion'
         )
+      case AIModel.CV2:
+        return renderModelDesc(
+          'OpenCV Image Inpainting',
+          'https://docs.opencv.org/4.6.0/df/d3d/tutorial_py_inpainting.html',
+          'https://docs.opencv.org/4.6.0/df/d3d/tutorial_py_inpainting.html'
+        )
       default:
         return <></>
     }

lama_cleaner/app/src/store/Atoms.tsx

@@ -10,6 +10,7 @@ export enum AIModel {
   MAT = 'mat',
   FCF = 'fcf',
   SD14 = 'sd1.4',
+  CV2 = 'cv2',
 }
 
 export const fileState = atom<File | undefined>({
@@ -132,6 +133,11 @@ export interface HDSettings {
 
 type ModelsHDSettings = { [key in AIModel]: HDSettings }
 
+export enum CV2Flag {
+  INPAINT_NS = 'INPAINT_NS',
+  INPAINT_TELEA = 'INPAINT_TELEA',
+}
+
 export interface Settings {
   show: boolean
   showCroper: boolean
@@ -158,6 +164,10 @@ export interface Settings {
   sdSeed: number
   sdSeedFixed: boolean // true: use sdSeed, false: random generate seed on backend
   sdNumSamples: number
+
+  // For OpenCV2
+  cv2Radius: number
+  cv2Flag: CV2Flag
 }
 
 const defaultHDSettings: ModelsHDSettings = {
@@ -203,6 +213,13 @@ const defaultHDSettings: ModelsHDSettings = {
     hdStrategyCropMargin: 128,
     enabled: true,
   },
+  [AIModel.CV2]: {
+    hdStrategy: HDStrategy.RESIZE,
+    hdStrategyResizeLimit: 1080,
+    hdStrategyCropTrigerSize: 512,
+    hdStrategyCropMargin: 128,
+    enabled: true,
+  },
 }
 
 export enum SDSampler {
@@ -240,6 +257,10 @@ export const settingStateDefault: Settings = {
   sdSeed: 42,
   sdSeedFixed: true,
   sdNumSamples: 1,
+
+  // CV2
+  cv2Radius: 5,
+  cv2Flag: CV2Flag.INPAINT_NS,
 }
 
 const localStorageEffect =

lama_cleaner/model/opencv2.py

@@ -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

lama_cleaner/model_manager.py

@@ -4,9 +4,10 @@ from lama_cleaner.model.ldm import LDM
 from lama_cleaner.model.mat import MAT
 from lama_cleaner.model.sd import SD14
 from lama_cleaner.model.zits import ZITS
+from lama_cleaner.model.opencv2 import OpenCV2
 from lama_cleaner.schema import Config
 
-models = {"lama": LaMa, "ldm": LDM, "zits": ZITS, "mat": MAT, "fcf": FcF, "sd1.4": SD14}
+models = {"lama": LaMa, "ldm": LDM, "zits": ZITS, "mat": MAT, "fcf": FcF, "sd1.4": SD14, "cv2": OpenCV2}
 
 
 class ModelManager:

lama_cleaner/parse_args.py

@@ -10,7 +10,7 @@ def parse_args():
     parser.add_argument(
         "--model",
         default="lama",
-        choices=["lama", "ldm", "zits", "mat", "fcf", "sd1.4"],
+        choices=["lama", "ldm", "zits", "mat", "fcf", "sd1.4", "cv2"],
     )
     parser.add_argument(
         "--hf_access_token",