refactor: split canvas_node into python modules and enhance matting UI

python/image_utils.py

@@ -0,0 +1,74 @@
+"""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"]