V 2.2.7 lama iopaint 去水印

2024-06-06 13:12:04 +08:00
parent 8b011856c2
commit f27ec08724
39 changed files with 2495 additions and 547 deletions
@@ -0,0 +1,2 @@
+@echo off
+pyinstaller -F --upx-dir="C:\\Users\\27698\\Desktop\\upx-4.2.4-win64\upx.exe" lama_inpaint.py
@@ -0,0 +1,173 @@
+import io
+import os
+import sys
+from typing import Union
+import cv2
+import torch
+import numpy as np
+from PIL import Image
+
+
+sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8")
+
+# 判断sys.argv 的长度，如果小于2，说明没有传入参数，设置初始参数
+# if len(sys.argv) < 2:
+#     sys.argv = [
+#         "C:/Users/27698/Desktop/LAITool/resources/scripts/lama/lama_inpaint.exe",
+#         "-l",
+#         "C:\\Users\\27698\\Desktop\\测试\\mjTest\\data\\mask\\temp\\1717508661218.png",
+#         "C:\\Users\\27698\\Desktop\\测试\\mjTest\\data\\mask\\mask_temp_1717508662659.png",
+#         "C:\\Users\\27698\\Desktop\\测试\\mjTest\\data\\mask\\temp\\1717508564042.png",
+#     ]
+print(sys.argv)
+
+if getattr(sys, "frozen", False):
+    cript_directory = os.path.dirname(sys.executable)
+elif __file__:
+    cript_directory = os.path.dirname(__file__)
+
+link_name = os.path.join(os.path.expanduser("~"), "big_lama.pt")
+cu_name = os.path.join(cript_directory, "model\\big-lama.pt")
+mode_pa = link_name
+
+if len(sys.argv) < 2:
+    # # 判断model_path是否存在，如果不存在，设置默认值
+    if not os.path.exists(link_name):
+        os.system(f'mklink "{link_name}" "{cu_name}"')
+    print("Params: <runtime-config.json>")
+    sys.exit(0)
+
+
+def get_image(image):
+    if isinstance(image, Image.Image):
+        img = np.array(image)
+    elif isinstance(image, np.ndarray):
+        img = image.copy()
+    else:
+        raise Exception("Input image should be either PIL Image or numpy array!")
+
+    if img.ndim == 3:
+        img = np.transpose(img, (2, 0, 1))  # chw
+    elif img.ndim == 2:
+        img = img[np.newaxis, ...]
+
+    assert img.ndim == 3
+
+    img = img.astype(np.float32) / 255
+    return img
+
+
+def ceil_modulo(x, mod):
+    if x % mod == 0:
+        return x
+    return (x // mod + 1) * mod
+
+
+def scale_image(img, factor, interpolation=cv2.INTER_AREA):
+    if img.shape[0] == 1:
+        img = img[0]
+    else:
+        img = np.transpose(img, (1, 2, 0))
+
+    img = cv2.resize(img, dsize=None, fx=factor, fy=factor, interpolation=interpolation)
+
+    if img.ndim == 2:
+        img = img[None, ...]
+    else:
+        img = np.transpose(img, (2, 0, 1))
+    return img
+
+
+def pad_img_to_modulo(img, mod):
+    channels, height, width = img.shape
+    out_height = ceil_modulo(height, mod)
+    out_width = ceil_modulo(width, mod)
+    return np.pad(
+        img,
+        ((0, 0), (0, out_height - height), (0, out_width - width)),
+        mode="symmetric",
+    )
+
+
+def prepare_img_and_mask(image, mask, device, pad_out_to_modulo=8, scale_factor=None):
+    out_image = get_image(image)
+    out_mask = get_image(mask)
+
+    if scale_factor is not None:
+        out_image = scale_image(out_image, 1)
+        out_mask = scale_image(out_mask, scale_factor, interpolation=cv2.INTER_NEAREST)
+
+    if pad_out_to_modulo is not None and pad_out_to_modulo > 1:
+        out_image = pad_img_to_modulo(out_image, pad_out_to_modulo)
+        out_mask = pad_img_to_modulo(out_mask, pad_out_to_modulo)
+
+    out_image = torch.from_numpy(out_image).unsqueeze(0).to(device)
+    out_mask = torch.from_numpy(out_mask).unsqueeze(0).to(device)
+
+    out_mask = (out_mask > 0) * 1
+
+    return out_image, out_mask
+
+
+class LamaInpaint:
+    def __init__(
+        self,
+        device,
+        model_path=None,
+    ) -> None:
+        if model_path is None:
+            model_path = os.path.join(cript_directory, "model\\big-lama.pt")
+
+        self.model = torch.jit.load(model_path, map_location=device)
+        self.model.eval()
+        self.model.to(device)
+        self.device = device
+
+    def run(
+        self,
+        image: Union[Image.Image, np.ndarray],
+        mask: Union[Image.Image, np.ndarray],
+    ):
+        if isinstance(image, np.ndarray):
+            orig_height, orig_width = image.shape[:2]
+        else:
+            orig_height, orig_width = np.array(image).shape[:2]
+
+        # image_width = image.shape[1]
+        # mask_width = mask.shape[1]
+        scale = image.width / mask.width
+        image, mask = prepare_img_and_mask(image, mask, self.device, 8, scale)
+        with torch.inference_mode():
+            inpainted = self.model(image, mask)
+            cur_res = inpainted[0].permute(1, 2, 0).detach().cpu().numpy()
+            cur_res = np.clip(cur_res * 255, 0, 255).astype("uint8")
+            cur_res = cur_res[:orig_height, :orig_width]
+            return cur_res
+
+
+try:
+    de = "cpu"
+    if torch.cuda.is_available():
+        de = "cuda"
+
+    lama = LamaInpaint(de, mode_pa)
+
+    image_path = sys.argv[2]
+    mask_path = sys.argv[3]
+    output_path = sys.argv[4]
+
+    # 若是没有传递mask_path，需要自己计算mask区域
+    # 使用Image.open打开图片
+    image = Image.open(image_path).convert("RGB")
+    mask = Image.open(mask_path).convert("L")
+
+    res = lama.run(image, mask)
+    # 将修复后的图片保存到本地
+    img = Image.fromarray(res)
+    # 使用 save 方法将图像保存到文件
+    img.save(output_path)
+    sys.exit(0)
+
+except Exception as e:
+    print(e)
+    sys.exit(str(e))
@@ -0,0 +1,43 @@
+# -*- mode: python ; coding: utf-8 -*-
+
+
+a = Analysis(
+    ['lama_inpaint.py'],
+    pathex=[],
+    binaries=[],
+    datas=[],
+    hiddenimports=[],
+    hookspath=[],
+    hooksconfig={},
+    runtime_hooks=[],
+    excludes=[],
+    noarchive=False,
+)
+pyz = PYZ(a.pure)
+
+exe = EXE(
+    pyz,
+    a.scripts,
+    [],
+    exclude_binaries=True,
+    name='lama_inpaint',
+    debug=False,
+    bootloader_ignore_signals=False,
+    strip=False,
+    upx=True,
+    console=True,
+    disable_windowed_traceback=False,
+    argv_emulation=False,
+    target_arch=None,
+    codesign_identity=None,
+    entitlements_file=None,
+)
+coll = COLLECT(
+    exe,
+    a.binaries,
+    a.datas,
+    strip=False,
+    upx=True,
+    upx_exclude=[],
+    name='lama_inpaint',
+)