65 lines
2.0 KiB
Python
65 lines
2.0 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_dirs = None):
|
|
if sub_dirs is None:
|
|
return CATEGORYNAMESPACE
|
|
else:
|
|
return "{}/utils".format(CATEGORYNAMESPACE)
|
|
|
|
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="JPEG", quality=80)
|
|
base64_image = base64.b64encode(img_bytes.getvalue()).decode('utf-8')
|
|
return base64_image
|