V 2.2.7 lama iopaint 去水印
This commit is contained in:
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -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',
|
||||
)
|
||||
@@ -0,0 +1,37 @@
|
||||
# -*- 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,
|
||||
a.binaries,
|
||||
a.datas,
|
||||
[],
|
||||
name='lama_inpaint',
|
||||
debug=False,
|
||||
bootloader_ignore_signals=False,
|
||||
strip=False,
|
||||
upx=True,
|
||||
upx_exclude=[],
|
||||
runtime_tmpdir=None,
|
||||
console=True,
|
||||
disable_windowed_traceback=False,
|
||||
argv_emulation=False,
|
||||
target_arch=None,
|
||||
codesign_identity=None,
|
||||
entitlements_file=None,
|
||||
)
|
||||
Reference in New Issue
Block a user