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

Merge pull request #62 from Sanster/add_MAT

Browse Source 
Add mat
This commit is contained in:
Qing 2022-08-24 21:44:57 +08:00 committed by GitHub
commit b6d71c4733
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
17 changed files with 2166 additions and 23 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

10
lama_cleaner/app/build/asset-manifest.json
View File

@ -1,17 +1,17 @@
{
"files": {
"main.css": "/static/css/main.eb627daf.chunk.css",
"main.js": "/static/js/main.5fe518c4.chunk.js",
"main.js": "/static/js/main.4e53814e.chunk.js",
"runtime-main.js": "/static/js/runtime-main.5e86ac81.js",
"static/js/2.3a92b865.chunk.js": "/static/js/2.3a92b865.chunk.js",
"static/js/2.cf2e9421.chunk.js": "/static/js/2.cf2e9421.chunk.js",
"index.html": "/index.html",
"static/js/2.3a92b865.chunk.js.LICENSE.txt": "/static/js/2.3a92b865.chunk.js.LICENSE.txt",
"static/js/2.cf2e9421.chunk.js.LICENSE.txt": "/static/js/2.cf2e9421.chunk.js.LICENSE.txt",
"static/media/_index.scss": "/static/media/WorkSans-SemiBold.1e98db4e.ttf"
},
"entrypoints": [
"static/js/runtime-main.5e86ac81.js",
"static/js/2.3a92b865.chunk.js",
"static/js/2.cf2e9421.chunk.js",
"static/css/main.eb627daf.chunk.css",
"static/js/main.5fe518c4.chunk.js"
"static/js/main.4e53814e.chunk.js"
]
}

2
lama_cleaner/app/build/index.html
View File

@ -1 +1 @@
<!doctype html><html lang="en"><head><meta charset="utf-8"/><meta name="viewport" content="width=device-width,initial-scale=1,maximum-scale=1,user-scalable=0"/><meta name="theme-color" content="#ffffff"/><title>lama-cleaner - Image inpainting powered by LaMa</title><link href="/static/css/main.eb627daf.chunk.css" rel="stylesheet"></head><body><noscript>You need to enable JavaScript to run this app.</noscript><div id="root"></div><script>!function(e){function r(r){for(var n,l,a=r[0],f=r[1],i=r[2],p=0,s=[];p<a.length;p++)l=a[p],Object.prototype.hasOwnProperty.call(o,l)&&o[l]&&s.push(o[l][0]),o[l]=0;for(n in f)Object.prototype.hasOwnProperty.call(f,n)&&(e[n]=f[n]);for(c&&c(r);s.length;)s.shift()();return u.push.apply(u,i||[]),t()}function t(){for(var e,r=0;r<u.length;r++){for(var t=u[r],n=!0,a=1;a<t.length;a++){var f=t[a];0!==o[f]&&(n=!1)}n&&(u.splice(r--,1),e=l(l.s=t[0]))}return e}var n={},o={1:0},u=[];function l(r){if(n[r])return n[r].exports;var t=n[r]={i:r,l:!1,exports:{}};return e[r].call(t.exports,t,t.exports,l),t.l=!0,t.exports}l.m=e,l.c=n,l.d=function(e,r,t){l.o(e,r)||Object.defineProperty(e,r,{enumerable:!0,get:t})},l.r=function(e){"undefined"!=typeof Symbol&&Symbol.toStringTag&&Object.defineProperty(e,Symbol.toStringTag,{value:"Module"}),Object.defineProperty(e,"__esModule",{value:!0})},l.t=function(e,r){if(1&r&&(e=l(e)),8&r)return e;if(4&r&&"object"==typeof e&&e&&e.__esModule)return e;var t=Object.create(null);if(l.r(t),Object.defineProperty(t,"default",{enumerable:!0,value:e}),2&r&&"string"!=typeof e)for(var n in e)l.d(t,n,function(r){return e[r]}.bind(null,n));return t},l.n=function(e){var r=e&&e.__esModule?function(){return e.default}:function(){return e};return l.d(r,"a",r),r},l.o=function(e,r){return Object.prototype.hasOwnProperty.call(e,r)},l.p="/";var a=this["webpackJsonplama-cleaner"]=this["webpackJsonplama-cleaner"]||[],f=a.push.bind(a);a.push=r,a=a.slice();for(var i=0;i<a.length;i++)r(a[i]);var c=f;t()}([])</script><script src="/static/js/2.3a92b865.chunk.js"></script><script src="/static/js/main.5fe518c4.chunk.js"></script></body></html>
<!doctype html><html lang="en"><head><meta charset="utf-8"/><meta name="viewport" content="width=device-width,initial-scale=1,maximum-scale=1,user-scalable=0"/><meta name="theme-color" content="#ffffff"/><title>lama-cleaner - Image inpainting powered by LaMa</title><link href="/static/css/main.eb627daf.chunk.css" rel="stylesheet"></head><body><noscript>You need to enable JavaScript to run this app.</noscript><div id="root"></div><script>!function(e){function r(r){for(var n,l,a=r[0],f=r[1],i=r[2],p=0,s=[];p<a.length;p++)l=a[p],Object.prototype.hasOwnProperty.call(o,l)&&o[l]&&s.push(o[l][0]),o[l]=0;for(n in f)Object.prototype.hasOwnProperty.call(f,n)&&(e[n]=f[n]);for(c&&c(r);s.length;)s.shift()();return u.push.apply(u,i||[]),t()}function t(){for(var e,r=0;r<u.length;r++){for(var t=u[r],n=!0,a=1;a<t.length;a++){var f=t[a];0!==o[f]&&(n=!1)}n&&(u.splice(r--,1),e=l(l.s=t[0]))}return e}var n={},o={1:0},u=[];function l(r){if(n[r])return n[r].exports;var t=n[r]={i:r,l:!1,exports:{}};return e[r].call(t.exports,t,t.exports,l),t.l=!0,t.exports}l.m=e,l.c=n,l.d=function(e,r,t){l.o(e,r)||Object.defineProperty(e,r,{enumerable:!0,get:t})},l.r=function(e){"undefined"!=typeof Symbol&&Symbol.toStringTag&&Object.defineProperty(e,Symbol.toStringTag,{value:"Module"}),Object.defineProperty(e,"__esModule",{value:!0})},l.t=function(e,r){if(1&r&&(e=l(e)),8&r)return e;if(4&r&&"object"==typeof e&&e&&e.__esModule)return e;var t=Object.create(null);if(l.r(t),Object.defineProperty(t,"default",{enumerable:!0,value:e}),2&r&&"string"!=typeof e)for(var n in e)l.d(t,n,function(r){return e[r]}.bind(null,n));return t},l.n=function(e){var r=e&&e.__esModule?function(){return e.default}:function(){return e};return l.d(r,"a",r),r},l.o=function(e,r){return Object.prototype.hasOwnProperty.call(e,r)},l.p="/";var a=this["webpackJsonplama-cleaner"]=this["webpackJsonplama-cleaner"]||[],f=a.push.bind(a);a.push=r,a=a.slice();for(var i=0;i<a.length;i++)r(a[i]);var c=f;t()}([])</script><script src="/static/js/2.cf2e9421.chunk.js"></script><script src="/static/js/main.4e53814e.chunk.js"></script></body></html>

2
lama_cleaner/app/build/static/js/2.3a92b865.chunk.js
View File

File diff suppressed because one or more lines are too long

2
lama_cleaner/app/build/static/js/2.cf2e9421.chunk.js Normal file
View File

File diff suppressed because one or more lines are too long

0
lama_cleaner/app/build/static/js/2.3a92b865.chunk.js.LICENSE.txt → lama_cleaner/app/build/static/js/2.cf2e9421.chunk.js.LICENSE.txt
View File

1
lama_cleaner/app/build/static/js/main.4e53814e.chunk.js Normal file
View File

File diff suppressed because one or more lines are too long

1
lama_cleaner/app/build/static/js/main.5fe518c4.chunk.js
View File

File diff suppressed because one or more lines are too long

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}",

28
lama_cleaner/model/base.py
View File

@ -120,10 +120,30 @@ class InpaintModel:
w = box_w + config.hd_strategy_crop_margin * 2
h = box_h + config.hd_strategy_crop_margin * 2
l = max(cx - w // 2, 0)
t = max(cy - h // 2, 0)
r = min(cx + w // 2, img_w)
b = min(cy + h // 2, img_h)
_l = cx - w // 2
_r = cx + w // 2
_t = cy - h // 2
_b = cy + h // 2
l = max(_l, 0)
r = min(_r, img_w)
t = max(_t, 0)
b = min(_b, img_h)
# try to get more context when crop around image edge
if _l < 0:
r += abs(_l)
if _r > img_w:
l -= (_r - img_w)
if _t < 0:
b += abs(_t)
if _b > img_h:
t -= (_b - img_h)
l = max(l, 0)
r = min(r, img_w)
t = max(t, 0)
b = min(b, img_h)
crop_img = image[t:b, l:r, :]
crop_mask = mask[t:b, l:r]

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",
)

2
setup.py
View File

@ -21,7 +21,7 @@ def load_requirements():
# https://setuptools.readthedocs.io/en/latest/setuptools.html#including-data-files
setuptools.setup(
name="lama-cleaner",
version="0.15.0",
version="0.16.0",
author="PanicByte",
author_email="cwq1913@gmail.com",
description="Image inpainting tool powered by SOTA AI Model",