114 lines
3.3 KiB
Python
114 lines
3.3 KiB
Python
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
def tensor_to_pil(tensor):
|
|
"""
|
|
Convert a PyTorch tensor to a PIL Image.
|
|
Uses ComfyUI's standard BHWC format with values in [0,1].
|
|
|
|
Args:
|
|
tensor: PyTorch tensor in BHWC format with values in [0,1]
|
|
|
|
Returns:
|
|
PIL.Image object
|
|
"""
|
|
# Handle different tensor shapes
|
|
if tensor.ndim == 4:
|
|
# BHWC format - take first image if batched
|
|
tensor = tensor[0]
|
|
|
|
# Convert to numpy and scale to [0, 255]
|
|
img_np = (tensor.cpu().numpy() * 255.0).astype(np.uint8)
|
|
|
|
# Create appropriate PIL image based on number of channels
|
|
if img_np.shape[2] == 4: # RGBA
|
|
img = Image.fromarray(img_np, 'RGBA')
|
|
elif img_np.shape[2] == 3: # RGB
|
|
img = Image.fromarray(img_np, 'RGB')
|
|
elif img_np.shape[2] == 1: # Grayscale
|
|
img = Image.fromarray(img_np.squeeze(2), 'L')
|
|
else:
|
|
raise ValueError(f"Unsupported tensor shape: {tensor.shape}")
|
|
|
|
return img
|
|
|
|
def pil_to_tensor(img):
|
|
"""
|
|
Convert a PIL Image to a PyTorch tensor.
|
|
Returns tensor in ComfyUI's standard BHWC format with float32 [0,1] range.
|
|
|
|
Args:
|
|
img: PIL Image
|
|
|
|
Returns:
|
|
PyTorch tensor in BHWC format with float32 [0,1] range
|
|
"""
|
|
# 添加类型检查以避免元组错误
|
|
if isinstance(img, tuple):
|
|
print("[警告] pil_to_tensor接收到元组而不是PIL图像，正在提取第一个元素")
|
|
if len(img) > 0 and hasattr(img[0], 'mode'):
|
|
img = img[0] # 尝试使用元组的第一个元素
|
|
else:
|
|
print("[错误] 无法从元组中提取有效图像")
|
|
# 创建一个1x1的黑色图像作为后备方案
|
|
img = Image.new('RGBA', (1, 1), (0, 0, 0, 0))
|
|
|
|
# 检查是否是PIL图像对象
|
|
if not hasattr(img, 'mode'):
|
|
print(f"[错误] 无效的图像对象类型: {type(img)}")
|
|
# 创建一个1x1的黑色图像作为后备方案
|
|
img = Image.new('RGBA', (1, 1), (0, 0, 0, 0))
|
|
|
|
# Convert to RGB or RGBA if needed
|
|
if img.mode not in ('RGB', 'RGBA'):
|
|
if 'A' in img.mode:
|
|
img = img.convert('RGBA')
|
|
else:
|
|
img = img.convert('RGB')
|
|
|
|
# Convert to numpy array and normalize to [0, 1]
|
|
img_np = np.array(img).astype(np.float32) / 255.0
|
|
|
|
# Add batch dimension for BHWC format
|
|
tensor = torch.from_numpy(img_np).unsqueeze(0)
|
|
|
|
return tensor
|
|
|
|
def create_alpha_mask(tensor):
|
|
"""
|
|
Create a single-channel alpha mask from an RGBA tensor.
|
|
|
|
Args:
|
|
tensor: PyTorch tensor in BHWC format with RGBA channels
|
|
|
|
Returns:
|
|
Alpha mask tensor in BHWC format with a single channel
|
|
"""
|
|
if tensor.shape[3] != 4:
|
|
return None
|
|
|
|
# Extract alpha channel and keep batch dimension
|
|
alpha = tensor[:, :, :, 3:4]
|
|
|
|
return alpha
|
|
|
|
def clean_alpha_mask(mask_tensor):
|
|
"""
|
|
Clean up alpha mask by removing any semi-transparent pixels.
|
|
|
|
Args:
|
|
mask_tensor: Alpha mask tensor in BHWC format (B, H, W, 1)
|
|
|
|
Returns:
|
|
Cleaned alpha mask tensor (B, H, W, 1)
|
|
"""
|
|
if mask_tensor is None:
|
|
return None
|
|
|
|
# Threshold the mask to make it binary
|
|
threshold = 0.05
|
|
cleaned_mask = (mask_tensor > threshold).float()
|
|
|
|
return cleaned_mask
|