kamp/IOPaint
1
0
Fork 0
Code Issues Pull Requests Releases 1 Activity

add MAT model

Browse Source
This commit is contained in:
Qing 2022-08-22 23:24:02 +08:00
parent a5e840765e
commit 6d2b24ed6b
8 changed files with 2132 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
README.md
View File

@ -26,6 +26,7 @@
1. [LaMa](https://github.com/saic-mdal/lama)
1. [LDM](https://github.com/CompVis/latent-diffusion)
1. [ZITS](https://github.com/DQiaole/ZITS_inpainting)
1. [MAT](https://github.com/fenglinglwb/MAT)
- Support CPU & GPU
- Various high-resolution image processing [strategy](#high-resolution-strategy)
- Run as a desktop APP
@ -36,7 +37,7 @@
| ---------------------- | --------------------------------------------- | --------------------------------------------------- |
| Remove unwanted things | ![unwant_object2](./assets/unwant_object.jpg) | ![unwant_object2](./assets/unwant_object_clean.jpg) |
| Remove unwanted person | ![unwant_person](./assets/unwant_person.jpg) | ![unwant_person](./assets/unwant_person_clean.jpg) |
| Remove Text | ![text](./assets/unwant_text.jpg) | ![watermark_clean](./assets/unwant_text_clean.jpg) |
| Remove Text | ![text](./assets/unwant_text.jpg) | ![text](./assets/unwant_text_clean.jpg) |
| Remove watermark | ![watermark](./assets/watermark.jpg) | ![watermark_clean](./assets/watermark_cleanup.jpg) |
| Fix old photo | ![oldphoto](./assets/old_photo.jpg) | ![oldphoto_clean](./assets/old_photo_clean.jpg) |
@ -69,6 +70,7 @@ Available arguments:
| LaMa | :+1: Generalizes well on high resolutions(~2k)<br/> | |
| LDM | :+1: Possiblablity to get better and more detail result <br/> :+1: The balance of time and quality can be achieved by adjusting `steps` <br/> :neutral_face: Slower than GAN model<br/> :neutral_face: Need more GPU memory | `Steps`: You can get better result with large steps, but it will be more time-consuming <br/> `Sampler`: ddim or [plms](https://arxiv.org/abs/2202.09778). In general plms can get [better results](https://github.com/Sanster/lama-cleaner/releases/tag/0.13.0) with fewer steps |
| ZITS | :+1: Better holistic structures compared with previous methods <br/> :neutral_face: Wireframe module is **very** slow on CPU | `Wireframe`: Enable edge and line detect |
| MAT | TODO | |
### LaMa vs LDM

8
lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx
View File

@ -131,6 +131,8 @@ function ModelSettingBlock() {
return renderLDMModelDesc()
case AIModel.ZITS:
return renderZITSModelDesc()
case AIModel.MAT:
return undefined
default:
return <></>
}
@ -156,6 +158,12 @@ function ModelSettingBlock() {
'https://arxiv.org/abs/2203.00867',
'https://github.com/DQiaole/ZITS_inpainting'
)
case AIModel.MAT:
return renderModelDesc(
'Mask-Aware Transformer for Large Hole Image Inpainting',
'https://arxiv.org/pdf/2203.15270.pdf',
'https://github.com/fenglinglwb/MAT'
)
default:
return <></>
}

7
lama_cleaner/app/src/store/Atoms.tsx
View File

@ -7,6 +7,7 @@ export enum AIModel {
LAMA = 'lama',
LDM = 'ldm',
ZITS = 'zits',
MAT = 'mat',
}
export const fileState = atom<File | undefined>({
@ -80,6 +81,12 @@ const defaultHDSettings: ModelsHDSettings = {
hdStrategyCropTrigerSize: 1024,
hdStrategyCropMargin: 128,
},
[AIModel.MAT]: {
hdStrategy: HDStrategy.CROP,
hdStrategyResizeLimit: 1024,
hdStrategyCropTrigerSize: 512,
hdStrategyCropMargin: 128,
},
}
export const settingStateDefault: Settings = {

20
lama_cleaner/helper.py
View File

@ -53,6 +53,26 @@ def load_jit_model(url_or_path, device):
return model
def load_model(model: torch.nn.Module, url_or_path, device):
if os.path.exists(url_or_path):
model_path = url_or_path
else:
model_path = download_model(url_or_path)
try:
state_dict = torch.load(model_path, map_location='cpu')
model.load_state_dict(state_dict, strict=True)
model.to(device)
logger.info(f"Load model from: {model_path}")
except:
logger.error(
f"Failed to load {model_path}, delete model and restart lama-cleaner"
)
exit(-1)
model.eval()
return model
def numpy_to_bytes(image_numpy: np.ndarray, ext: str) -> bytes:
data = cv2.imencode(
f".{ext}",

2064
lama_cleaner/model/mat.py Normal file
View File

File diff suppressed because it is too large Load Diff

4
lama_cleaner/model_manager.py
View File

@ -1,12 +1,14 @@
from lama_cleaner.model.lama import LaMa
from lama_cleaner.model.ldm import LDM
from lama_cleaner.model.mat import MAT
from lama_cleaner.model.zits import ZITS
from lama_cleaner.schema import Config
models = {
'lama': LaMa,
'ldm': LDM,
'zits': ZITS
'zits': ZITS,
'mat': MAT
}

2
lama_cleaner/parse_args.py
View File

@ -7,7 +7,7 @@ def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--host", default="127.0.0.1")
parser.add_argument("--port", default=8080, type=int)
parser.add_argument("--model", default="lama", choices=["lama", "ldm", "zits"])
parser.add_argument("--model", default="lama", choices=["lama", "ldm", "zits", "mat"])
parser.add_argument("--device", default="cuda", type=str, choices=["cuda", "cpu"])
parser.add_argument("--gui", action="store_true", help="Launch as desktop app")
parser.add_argument(

32
lama_cleaner/tests/test_model.py
View File

@ -11,13 +11,19 @@ from lama_cleaner.schema import Config, HDStrategy, LDMSampler
current_dir = Path(__file__).parent.absolute().resolve()
def get_data(fx=1):
def get_data(fx=1, fy=1.0):
img = cv2.imread(str(current_dir / "image.png"))
img = cv2.cvtColor(img, cv2.COLOR_BGRA2RGB)
mask = cv2.imread(str(current_dir / "mask.png"), cv2.IMREAD_GRAYSCALE)
# img = cv2.imread("/Users/qing/code/github/MAT/test_sets/Places/images/test1.jpg")
# img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# mask = cv2.imread("/Users/qing/code/github/MAT/test_sets/Places/masks/mask1.png", cv2.IMREAD_GRAYSCALE)
# mask = 255 - mask
if fx != 1:
img = cv2.resize(img, None, fx=fx, fy=1)
mask = cv2.resize(mask, None, fx=fx, fy=1)
img = cv2.resize(img, None, fx=fx, fy=fy)
mask = cv2.resize(mask, None, fx=fx, fy=fy)
return img, mask
@ -34,8 +40,8 @@ def get_config(strategy, **kwargs):
return Config(**data)
def assert_equal(model, config, gt_name, fx=1):
img, mask = get_data(fx=fx)
def assert_equal(model, config, gt_name, fx=1, fy=1):
img, mask = get_data(fx=fx, fy=fy)
res = model(img, mask, config)
cv2.imwrite(
str(current_dir / gt_name),
@ -111,6 +117,20 @@ def test_zits(strategy, zits_wireframe):
assert_equal(
model,
cfg,
f"zits_{strategy[0].upper() + strategy[1:]}_wireframe_{zits_wireframe}_fx_{fx}_result.png",
f"zits_{strategy.capitalize()}_wireframe_{zits_wireframe}_fx_{fx}_result.png",
fx=fx,
)
@pytest.mark.parametrize(
"strategy", [HDStrategy.ORIGINAL]
)
def test_mat(strategy):
model = ModelManager(name="mat", device="cpu")
cfg = get_config(strategy)
assert_equal(
model,
cfg,
f"mat_{strategy.capitalize()}_result.png",
)