Add OpenCV2 model
This commit is contained in:
parent
bc98ea256a
commit
35cbbd653c
@ -1,6 +1,6 @@
|
|||||||
import React, { ReactNode } from 'react'
|
import React, { ReactNode } from 'react'
|
||||||
import { useRecoilState } from 'recoil'
|
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 Selector from '../shared/Selector'
|
||||||
import { Switch, SwitchThumb } from '../shared/Switch'
|
import { Switch, SwitchThumb } from '../shared/Switch'
|
||||||
import Tooltip from '../shared/Tooltip'
|
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 => {
|
const renderOptionDesc = (): ReactNode => {
|
||||||
switch (setting.model) {
|
switch (setting.model) {
|
||||||
case AIModel.LAMA:
|
case AIModel.LAMA:
|
||||||
@ -147,6 +183,8 @@ function ModelSettingBlock() {
|
|||||||
return renderFCFModelDesc()
|
return renderFCFModelDesc()
|
||||||
case AIModel.SD14:
|
case AIModel.SD14:
|
||||||
return undefined
|
return undefined
|
||||||
|
case AIModel.CV2:
|
||||||
|
return renderOpenCV2Desc()
|
||||||
default:
|
default:
|
||||||
return <></>
|
return <></>
|
||||||
}
|
}
|
||||||
@ -190,6 +228,12 @@ function ModelSettingBlock() {
|
|||||||
'https://ommer-lab.com/research/latent-diffusion-models/',
|
'https://ommer-lab.com/research/latent-diffusion-models/',
|
||||||
'https://github.com/CompVis/stable-diffusion'
|
'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:
|
default:
|
||||||
return <></>
|
return <></>
|
||||||
}
|
}
|
||||||
|
@ -10,6 +10,7 @@ export enum AIModel {
|
|||||||
MAT = 'mat',
|
MAT = 'mat',
|
||||||
FCF = 'fcf',
|
FCF = 'fcf',
|
||||||
SD14 = 'sd1.4',
|
SD14 = 'sd1.4',
|
||||||
|
CV2 = 'cv2',
|
||||||
}
|
}
|
||||||
|
|
||||||
export const fileState = atom<File | undefined>({
|
export const fileState = atom<File | undefined>({
|
||||||
@ -132,6 +133,11 @@ export interface HDSettings {
|
|||||||
|
|
||||||
type ModelsHDSettings = { [key in AIModel]: HDSettings }
|
type ModelsHDSettings = { [key in AIModel]: HDSettings }
|
||||||
|
|
||||||
|
export enum CV2Flag {
|
||||||
|
INPAINT_NS = 'INPAINT_NS',
|
||||||
|
INPAINT_TELEA = 'INPAINT_TELEA',
|
||||||
|
}
|
||||||
|
|
||||||
export interface Settings {
|
export interface Settings {
|
||||||
show: boolean
|
show: boolean
|
||||||
showCroper: boolean
|
showCroper: boolean
|
||||||
@ -158,6 +164,10 @@ export interface Settings {
|
|||||||
sdSeed: number
|
sdSeed: number
|
||||||
sdSeedFixed: boolean // true: use sdSeed, false: random generate seed on backend
|
sdSeedFixed: boolean // true: use sdSeed, false: random generate seed on backend
|
||||||
sdNumSamples: number
|
sdNumSamples: number
|
||||||
|
|
||||||
|
// For OpenCV2
|
||||||
|
cv2Radius: number
|
||||||
|
cv2Flag: CV2Flag
|
||||||
}
|
}
|
||||||
|
|
||||||
const defaultHDSettings: ModelsHDSettings = {
|
const defaultHDSettings: ModelsHDSettings = {
|
||||||
@ -203,6 +213,13 @@ const defaultHDSettings: ModelsHDSettings = {
|
|||||||
hdStrategyCropMargin: 128,
|
hdStrategyCropMargin: 128,
|
||||||
enabled: true,
|
enabled: true,
|
||||||
},
|
},
|
||||||
|
[AIModel.CV2]: {
|
||||||
|
hdStrategy: HDStrategy.RESIZE,
|
||||||
|
hdStrategyResizeLimit: 1080,
|
||||||
|
hdStrategyCropTrigerSize: 512,
|
||||||
|
hdStrategyCropMargin: 128,
|
||||||
|
enabled: true,
|
||||||
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
export enum SDSampler {
|
export enum SDSampler {
|
||||||
@ -240,6 +257,10 @@ export const settingStateDefault: Settings = {
|
|||||||
sdSeed: 42,
|
sdSeed: 42,
|
||||||
sdSeedFixed: true,
|
sdSeedFixed: true,
|
||||||
sdNumSamples: 1,
|
sdNumSamples: 1,
|
||||||
|
|
||||||
|
// CV2
|
||||||
|
cv2Radius: 5,
|
||||||
|
cv2Flag: CV2Flag.INPAINT_NS,
|
||||||
}
|
}
|
||||||
|
|
||||||
const localStorageEffect =
|
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.mat import MAT
|
||||||
from lama_cleaner.model.sd import SD14
|
from lama_cleaner.model.sd import SD14
|
||||||
from lama_cleaner.model.zits import ZITS
|
from lama_cleaner.model.zits import ZITS
|
||||||
|
from lama_cleaner.model.opencv2 import OpenCV2
|
||||||
from lama_cleaner.schema import Config
|
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:
|
class ModelManager:
|
||||||
|
@ -10,7 +10,7 @@ def parse_args():
|
|||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--model",
|
"--model",
|
||||||
default="lama",
|
default="lama",
|
||||||
choices=["lama", "ldm", "zits", "mat", "fcf", "sd1.4"],
|
choices=["lama", "ldm", "zits", "mat", "fcf", "sd1.4", "cv2"],
|
||||||
)
|
)
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--hf_access_token",
|
"--hf_access_token",
|
||||||
|
Loading…
Reference in New Issue
Block a user