75 lines
2.5 KiB
Python
75 lines
2.5 KiB
Python
"""Image and tensor conversion helpers shared by LayerForge endpoints."""
|
|
|
|
import base64
|
|
import io
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
from torchvision import transforms
|
|
|
|
from .node import log
|
|
|
|
|
|
def convert_base64_to_tensor(base64_str):
|
|
"""Convert a data-URL image into a BCHW tensor and optional alpha tensor."""
|
|
try:
|
|
img_data = base64.b64decode(base64_str.split(",")[1])
|
|
img = Image.open(io.BytesIO(img_data))
|
|
|
|
has_alpha = img.mode == "RGBA"
|
|
alpha = None
|
|
if has_alpha:
|
|
alpha = img.split()[3]
|
|
background = Image.new("RGB", img.size, (255, 255, 255))
|
|
background.paste(img, mask=alpha)
|
|
img = background
|
|
elif img.mode != "RGB":
|
|
img = img.convert("RGB")
|
|
|
|
img_tensor = transforms.ToTensor()(img).unsqueeze(0)
|
|
if has_alpha:
|
|
alpha_tensor = transforms.ToTensor()(alpha).unsqueeze(0)
|
|
return img_tensor, alpha_tensor
|
|
|
|
return img_tensor, None
|
|
except Exception as error:
|
|
log.error(f"Error in convert_base64_to_tensor: {error}")
|
|
raise
|
|
|
|
|
|
def convert_tensor_to_base64(tensor, alpha_mask=None, original_alpha=None):
|
|
"""Convert a tensor to a PNG data URL, optionally preserving alpha."""
|
|
try:
|
|
tensor = tensor.cpu()
|
|
|
|
if tensor.dim() == 4:
|
|
tensor = tensor.squeeze(0)
|
|
if tensor.dim() == 3 and tensor.shape[0] in [1, 3]:
|
|
tensor = tensor.permute(1, 2, 0)
|
|
|
|
img_array = (tensor.numpy() * 255).astype(np.uint8)
|
|
|
|
if alpha_mask is not None and original_alpha is not None:
|
|
alpha_mask = (alpha_mask.cpu().squeeze().numpy() * 255).astype(np.uint8)
|
|
original_alpha = (original_alpha.cpu().squeeze().numpy() * 255).astype(np.uint8)
|
|
combined_alpha = np.minimum(alpha_mask, original_alpha)
|
|
|
|
img = Image.fromarray(img_array, mode="RGB")
|
|
img.putalpha(Image.fromarray(combined_alpha, mode="L"))
|
|
elif img_array.shape[-1] == 1:
|
|
img = Image.fromarray(img_array.squeeze(-1), mode="L")
|
|
else:
|
|
img = Image.fromarray(img_array, mode="RGB")
|
|
|
|
buffer = io.BytesIO()
|
|
img.save(buffer, format="PNG")
|
|
img_str = base64.b64encode(buffer.getvalue()).decode()
|
|
return f"data:image/png;base64,{img_str}"
|
|
except Exception as error:
|
|
log.error(f"Error in convert_tensor_to_base64: {error}")
|
|
log.debug(f"Tensor shape: {tensor.shape}, dtype: {tensor.dtype}")
|
|
raise
|
|
|
|
|
|
__all__ = ["convert_base64_to_tensor", "convert_tensor_to_base64"]
|