Merge pull request #62 from Sanster/add_MAT

Add mat

README.md

@@ -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
 

lama_cleaner/app/build/asset-manifest.json

@@ -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"
   ]
 }

lama_cleaner/app/build/index.html

@@ -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>

lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx

@@ -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 <></>
     }

lama_cleaner/app/src/store/Atoms.tsx

@@ -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 = {

lama_cleaner/helper.py

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

lama_cleaner/model/base.py

@@ -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)
+        _l = cx - w // 2
-        t = max(cy - h // 2, 0)
+        _r = cx + w // 2
-        r = min(cx + w // 2, img_w)
+        _t = cy - h // 2
-        b = min(cy + h // 2, img_h)
+        _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]

lama_cleaner/model_manager.py

@@ -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
 }
 
 

lama_cleaner/parse_args.py

@@ -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(

lama_cleaner/tests/test_model.py

@@ -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)
+        img = cv2.resize(img, None, fx=fx, fy=fy)
-        mask = cv2.resize(mask, None, fx=fx, fy=1)
+        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):
+def assert_equal(model, config, gt_name, fx=1, fy=1):
-    img, mask = get_data(fx=fx)
+    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",
+    )

setup.py

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