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

0.27.1

Browse Source
This commit is contained in:
Qing 2022-11-18 23:07:02 +08:00
parent 4686c7cb78
commit c02d918368
2 changed files with 12 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
lama_cleaner/model/manga.py
View File

@ -1,4 +1,5 @@
import os import os
import random
import cv2 import cv2
import numpy as np import numpy as np
@ -80,6 +81,7 @@ class Manga(InpaintModel):
def init_model(self, device, **kwargs): def init_model(self, device, **kwargs):
self.inpaintor_model = load_jit_model(MANGA_INPAINTOR_MODEL_URL, device) self.inpaintor_model = load_jit_model(MANGA_INPAINTOR_MODEL_URL, device)
self.line_model = load_jit_model(MANGA_LINE_MODEL_URL, device) self.line_model = load_jit_model(MANGA_LINE_MODEL_URL, device)
self.seed = 42
@staticmethod @staticmethod
def is_downloaded() -> bool: def is_downloaded() -> bool:
@ -95,10 +97,17 @@ class Manga(InpaintModel):
mask: [H, W, 1] mask: [H, W, 1]
return: BGR IMAGE return: BGR IMAGE
""" """
seed = self.seed
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
gray_img = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) gray_img = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
gray_img = torch.from_numpy(gray_img[np.newaxis, np.newaxis, :, :].astype(np.float32)).to(self.device) gray_img = torch.from_numpy(gray_img[np.newaxis, np.newaxis, :, :].astype(np.float32)).to(self.device)
start = time.time() start = time.time()
lines = self.line_model(gray_img) lines = self.line_model(gray_img)
torch.cuda.empty_cache()
lines = torch.clamp(lines, 0, 255) lines = torch.clamp(lines, 0, 255)
logger.info(f"erika_model time: {time.time() - start}") logger.info(f"erika_model time: {time.time() - start}")
@ -106,12 +115,13 @@ class Manga(InpaintModel):
mask = mask.permute(0, 3, 1, 2) mask = mask.permute(0, 3, 1, 2)
mask = torch.where(mask > 0.5, 1.0, 0.0) mask = torch.where(mask > 0.5, 1.0, 0.0)
noise = torch.randn_like(mask) noise = torch.randn_like(mask)
ones = torch.ones_like(mask)
gray_img = gray_img / 255 * 2 - 1.0 gray_img = gray_img / 255 * 2 - 1.0
lines = lines / 255 * 2 - 1.0 lines = lines / 255 * 2 - 1.0
start = time.time() start = time.time()
inpainted_image = self.inpaintor_model(gray_img, lines, mask, noise) inpainted_image = self.inpaintor_model(gray_img, lines, mask, noise, ones)
logger.info(f"image_inpaintor_model time: {time.time() - start}") logger.info(f"image_inpaintor_model time: {time.time() - start}")
cur_res = inpainted_image[0].permute(1, 2, 0).detach().cpu().numpy() cur_res = inpainted_image[0].permute(1, 2, 0).detach().cpu().numpy()

2
setup.py
View File

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