85 lines
2.7 KiB
Python
85 lines
2.7 KiB
Python
import base64
|
|
from io import BytesIO
|
|
from PIL import Image
|
|
import numpy as np
|
|
import torch
|
|
import cv2
|
|
|
|
CATEGORYNAMESPACE = 'OpenOutpaint-Serving'
|
|
|
|
def get_category(sub_dir = None):
|
|
if sub_dir is None:
|
|
return CATEGORYNAMESPACE
|
|
else:
|
|
return f"{CATEGORYNAMESPACE}/{sub_dir}"
|
|
|
|
def _strip_prefix(s: str, prefix: str) -> str:
|
|
if s.startswith(prefix):
|
|
return s[len(prefix) :]
|
|
return s
|
|
|
|
|
|
def convert_color(image):
|
|
if len(image.shape) > 2 and image.shape[2] >= 4:
|
|
return cv2.cvtColor(image, cv2.COLOR_BGRA2RGB)
|
|
return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
|
|
def base64_to_image(base64_str):
|
|
base64_str = _strip_prefix(base64_str, "data:image/png;base64,")
|
|
nparr = np.frombuffer(base64.b64decode(base64_str), np.uint8)
|
|
result = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED)
|
|
result = convert_color(result)
|
|
result = result.astype(np.float32) / 255.0
|
|
image = torch.from_numpy(result)[None,]
|
|
return image
|
|
|
|
def base64_to_mask(base64_str):
|
|
base64_str = _strip_prefix(base64_str, "data:image/png;base64,")
|
|
nparr = np.frombuffer(base64.b64decode(base64_str), np.uint8)
|
|
result = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED)
|
|
result = result.astype(np.float32) / 255.0
|
|
image = torch.from_numpy(result)
|
|
if image.dim() == 3: # RGB(A) input, use red channel
|
|
image = image[:, :, 0]
|
|
return image.unsqueeze(0)
|
|
|
|
def image_to_base64(image):
|
|
img_np = (image.cpu().numpy() * 255).astype('uint8')
|
|
img_bytes = BytesIO()
|
|
Image.fromarray(img_np.squeeze()).save(img_bytes, format='PNG')
|
|
base64_image = base64.b64encode(img_bytes.getvalue()).decode('utf-8')
|
|
return base64_image
|
|
|
|
def images_to_base64(images):
|
|
base64_images = []
|
|
for image in images:
|
|
base64_image = image_to_base64(image)
|
|
base64_images.append(base64_image)
|
|
return base64_images
|
|
|
|
def preview_to_base64(image):
|
|
img_bytes = BytesIO()
|
|
image.save(img_bytes, format="PNG", quality=95)
|
|
base64_image = base64.b64encode(img_bytes.getvalue()).decode('utf-8')
|
|
return base64_image
|
|
|
|
import json
|
|
def print_list_or_dic(header_str, data, test_json = False):
|
|
output = ""
|
|
output += f"\n========== {header_str}: start ==========\n"
|
|
for key, value in enumerate(data) if type(data) is list else data.items():
|
|
if key in ["images", "image", "mask", "init_images"]:
|
|
truncated_value = str(value)[:30]
|
|
else:
|
|
truncated_value = str(value)[:500]
|
|
output += f"> ({type(value)}) {key}: {truncated_value}\n"
|
|
|
|
if test_json:
|
|
try:
|
|
j = json.dumps(data).encode('utf-8')
|
|
except Exception as e:
|
|
output += f"ERROR: {e}\ndata: <{type(data)}> {data}\n"
|
|
|
|
output += f"========== {header_str}: end ==========\n"
|
|
print(output)
|