File size: 5,242 Bytes
e041d7d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 |
import base64
import io
import json
from pathlib import Path
from typing import Dict, Optional
import cv2
import psutil
from PIL import Image
from loguru import logger
from rich.console import Console
from rich.progress import (
Progress,
SpinnerColumn,
TimeElapsedColumn,
MofNCompleteColumn,
TextColumn,
BarColumn,
TaskProgressColumn,
)
from iopaint.helper import pil_to_bytes_single
from iopaint.model.utils import torch_gc
from iopaint.model_manager import ModelManager
from iopaint.schema import InpaintRequest
import numpy as np
def glob_images(path: Path) -> Dict[str, Path]:
# png/jpg/jpeg
if path.is_file():
return {path.stem: path}
elif path.is_dir():
res = {}
for it in path.glob("*.*"):
if it.suffix.lower() in [".png", ".jpg", ".jpeg"]:
res[it.stem] = it
return res
# def batch_inpaint(
# model: str,
# device,
# image: Path,
# mask: Path,
# config: Optional[Path] = None,
# concat: bool = False,
# ):
# if config is None:
# inpaint_request = InpaintRequest()
# else:
# with open(config, "r", encoding="utf-8") as f:
# inpaint_request = InpaintRequest(**json.load(f))
#
# model_manager = ModelManager(name=model, device=device)
#
# img = cv2.imread(str(image))
# img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
#
# mask_img = cv2.imread(str(mask), cv2.IMREAD_GRAYSCALE)
#
# if mask_img.shape[:2] != img.shape[:2]:
# mask_img = cv2.resize(
# mask_img,
# (img.shape[1], img.shape[0]),
# interpolation=cv2.INTER_NEAREST,
# )
#
# mask_img[mask_img >= 127] = 255
# mask_img[mask_img < 127] = 0
#
# # bgr
# inpaint_result = model_manager(img, mask_img, inpaint_request)
# inpaint_result = cv2.cvtColor(inpaint_result, cv2.COLOR_BGR2RGB)
#
# if concat:
# mask_img = cv2.cvtColor(mask_img, cv2.COLOR_GRAY2RGB)
# inpaint_result = cv2.hconcat([img, mask_img, inpaint_result])
#
# # Convert the NumPy array to PIL Image
# pil_image = Image.fromarray(inpaint_result)
#
# # Encode the PIL Image as base64 string
# with io.BytesIO() as output_buffer:
# pil_image.save(output_buffer, format='PNG')
# base64_image = base64.b64encode(output_buffer.getvalue()).decode('utf-8')
#
# return base64_image
def batch_inpaint(
model: str,
device,
input_base64: str,
mask_base64: str,
config_base64: Optional[str] = None,
concat: bool = False,
):
if config_base64 is None:
inpaint_request = InpaintRequest()
else:
config_json = base64.b64decode(config_base64)
inpaint_request = InpaintRequest(**json.loads(config_json))
model_manager = ModelManager(name=model, device=device)
# Decode input image from base64
input_image_data = base64.b64decode(input_base64)
input_image = cv2.imdecode(np.frombuffer(input_image_data, np.uint8), cv2.IMREAD_COLOR)
# Decode mask image from base64
mask_image_data = base64.b64decode(mask_base64)
mask_image = cv2.imdecode(np.frombuffer(mask_image_data, np.uint8), cv2.IMREAD_GRAYSCALE)
if mask_image.shape[:2] != input_image.shape[:2]:
mask_image = cv2.resize(
mask_image,
(input_image.shape[1], input_image.shape[0]),
interpolation=cv2.INTER_NEAREST,
)
mask_image[mask_image >= 127] = 255
mask_image[mask_image < 127] = 0
# Run inpainting
inpaint_result = model_manager(input_image, mask_image, inpaint_request)
if concat:
mask_image = cv2.cvtColor(mask_image, cv2.COLOR_GRAY2RGB)
inpaint_result = cv2.hconcat([input_image, mask_image, inpaint_result])
# Convert NumPy array to PIL Image
pil_image = Image.fromarray(inpaint_result)
# Encode PIL Image to base64 string
with io.BytesIO() as output_buffer:
pil_image.save(output_buffer, format='PNG')
base64_image = base64.b64encode(output_buffer.getvalue()).decode('utf-8')
return base64_image
def batch_inpaint_cv2(
model: str,
device,
input_base: str,
mask_base: str,
config_base64: Optional[str] = None,
concat: bool = False,
):
if config_base64 is None:
inpaint_request = InpaintRequest()
else:
config_json = base64.b64decode(config_base64)
inpaint_request = InpaintRequest(**json.loads(config_json))
model_manager = ModelManager(name=model, device=device)
# Decode input image from base
input_image = input_base
# Decode mask image from base
mask_image = mask_base
if mask_image.shape[:2] != input_image.shape[:2]:
mask_image = cv2.resize(
mask_image,
(input_image.shape[1], input_image.shape[0]),
interpolation=cv2.INTER_NEAREST,
)
mask_image[mask_image >= 127] = 255
mask_image[mask_image < 127] = 0
# Run inpainting
inpaint_result = model_manager(input_image, mask_image, inpaint_request)
if concat:
mask_image = cv2.cvtColor(mask_image, cv2.COLOR_GRAY2RGB)
inpaint_result = cv2.hconcat([input_image, mask_image, inpaint_result])
return inpaint_result |