Add OpenCV2 model
This commit is contained in:
parent
bc98ea256a
commit
35cbbd653c
@ -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 <></>
|
||||
}
|
||||
|
@ -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 =
|
||||
|
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
|
@ -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:
|
||||
|
@ -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",
|
||||
|
Loading…
Reference in New Issue
Block a user