68 lines
2.1 KiB
Python
68 lines
2.1 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
from io import BytesIO
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
|
|
|
|
def _iter_image_tensors(image: torch.Tensor) -> list[torch.Tensor]:
|
|
if image.ndim == 4:
|
|
return [image[index] for index in range(image.shape[0])]
|
|
if image.ndim == 3:
|
|
return [image]
|
|
raise ValueError("Expected image tensor with shape [B, H, W, C] or [H, W, C].")
|
|
|
|
|
|
def _normalize_image_tensor(image: torch.Tensor) -> torch.Tensor:
|
|
if image.ndim == 4:
|
|
image = image[0]
|
|
if image.ndim != 3:
|
|
raise ValueError("Expected image tensor with shape [B, H, W, C] or [H, W, C].")
|
|
return image.clamp(0.0, 1.0).cpu()
|
|
|
|
|
|
def tensor_to_png_bytes(image: torch.Tensor) -> bytes:
|
|
normalized = _normalize_image_tensor(image)
|
|
image_array = (normalized.numpy() * 255).astype(np.uint8)
|
|
pil_image = Image.fromarray(image_array)
|
|
buffer = BytesIO()
|
|
pil_image.save(buffer, format="PNG")
|
|
return buffer.getvalue()
|
|
|
|
|
|
def tensor_to_base64_png(image: torch.Tensor) -> str:
|
|
raw = tensor_to_png_bytes(image)
|
|
encoded = base64.b64encode(raw).decode("utf-8")
|
|
return f"data:image/png;base64,{encoded}"
|
|
|
|
|
|
def tensor_batch_to_base64_pngs(image: torch.Tensor) -> list[str]:
|
|
return [tensor_to_base64_png(item) for item in _iter_image_tensors(image)]
|
|
|
|
|
|
def base64_data_uri_to_bytes(data_uri: str) -> bytes:
|
|
if "," not in data_uri:
|
|
raise ValueError("Expected a valid data URI.")
|
|
return base64.b64decode(data_uri.split(",", 1)[1])
|
|
|
|
|
|
def decode_image_bytes(raw_bytes: bytes) -> torch.Tensor:
|
|
pil_image = Image.open(BytesIO(raw_bytes)).convert("RGB")
|
|
image_array = np.asarray(pil_image).astype(np.float32) / 255.0
|
|
return torch.from_numpy(image_array).unsqueeze(0)
|
|
|
|
|
|
def stack_image_tensors(images: list[torch.Tensor]) -> torch.Tensor:
|
|
if not images:
|
|
raise ValueError("Expected at least one image tensor to stack.")
|
|
normalized = []
|
|
for image in images:
|
|
if image.ndim == 3:
|
|
normalized.append(image.unsqueeze(0))
|
|
else:
|
|
normalized.append(image)
|
|
return torch.cat(normalized, dim=0)
|