diff --git a/lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx b/lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx index 19aeea0..41a681f 100644 --- a/lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx +++ b/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 ( + <> + { + const val = value.length === 0 ? 0 : parseInt(value, 10) + setSettingState(old => { + return { ...old, cv2Radius: 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 <> } diff --git a/lama_cleaner/app/src/store/Atoms.tsx b/lama_cleaner/app/src/store/Atoms.tsx index 9d8b985..85f842a 100644 --- a/lama_cleaner/app/src/store/Atoms.tsx +++ b/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({ @@ -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 = diff --git a/lama_cleaner/model/opencv2.py b/lama_cleaner/model/opencv2.py new file mode 100644 index 0000000..1a69baa --- /dev/null +++ b/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 diff --git a/lama_cleaner/model_manager.py b/lama_cleaner/model_manager.py index 1419183..ea623a4 100644 --- a/lama_cleaner/model_manager.py +++ b/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: diff --git a/lama_cleaner/parse_args.py b/lama_cleaner/parse_args.py index 3b0218d..f4b720e 100644 --- a/lama_cleaner/parse_args.py +++ b/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",