ADD 增加两个图像编辑节点

utils/image_utils.py

@@ -0,0 +1,73 @@
+import base64
+import io
+
+import torch
+from PIL import Image
+from torchvision import transforms
+
+
+def base64_to_tensor(base64_data: str) -> torch.Tensor:
+    """
+    将"data:image/xxx;base64,xxx"格式的图像数据转换为PyTorch张量
+
+    参数:
+        base64_data: 图像的Base64编码字符串
+
+    返回:
+        torch.Tensor: 形状为[C, H, W]的张量，取值范围为[0, 1]
+    """
+    # 分离数据前缀和实际Base64编码部分
+    if ';base64,' in base64_data:
+        _, encoded = base64_data.split(';base64,', 1)
+    else:
+        encoded = base64_data  # 假设直接提供了Base64编码部分
+
+    # 解码Base64数据
+    decoded_data = base64.b64decode(encoded)
+
+    # 使用PIL打开图像
+    image = Image.open(io.BytesIO(decoded_data))
+
+    # 转换为RGB模式（处理PNG的Alpha通道和WebP格式）
+    if image.mode != 'RGB':
+        image = image.convert('RGB')
+
+    # 转换为PyTorch张量
+    from torchvision import transforms
+    transform = transforms.Compose([
+        transforms.ToTensor()  # [H, W, C] -> [C, H, W]，并归一化到[0, 1]
+    ])
+    tensor = transform(image)
+
+    return tensor.unsqueeze(0).permute(0, 2, 3, 1)
+
+
+def tensor_to_image_bytes(tensor: torch.Tensor, format: str = 'PNG') -> bytes:
+    """
+    将PyTorch张量转换为图像字节流
+
+    参数:
+        tensor: 形状为[C, H, W]的图像张量，取值范围为[0, 1]
+        format: 图像格式，可选'PNG'、'JPEG'等
+
+    返回:
+        bytes: 图像的字节流数据
+    """
+    if tensor.dim() == 4:
+        if tensor.shape[0] > 1:
+            print("警告：输入张量包含多个图像，仅使用第一个")
+        tensor = tensor[0]  # 取批量中的第一张图像
+    tensor = tensor.permute(2, 0, 1)
+    # 确保张量在[0, 255]范围内
+    if tensor.max() <= 1.0:
+        tensor = tensor * 255
+
+    # 转换为PIL图像
+    image = transforms.ToPILImage()(tensor.byte())
+
+    # 保存为字节流
+    buffer = io.BytesIO()
+    image.save(buffer, format=format)
+    buffer.seek(0)  # 重置指针到开始位置
+
+    return buffer.getvalue()