| |
| import numpy as np |
| import torch |
| from PIL import Image |
| from comfy.utils import common_upscale |
|
|
|
|
| def tensor2pil(image: torch.Tensor) -> Image.Image: |
| return Image.fromarray(np.clip(255.0 * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) |
|
|
|
|
| def pil2tensor(image: Image.Image) -> torch.Tensor: |
| return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) |
|
|
|
|
| def pil2mask(image: Image.Image) -> torch.Tensor: |
| return torch.from_numpy(np.array(image.convert("L")).astype(np.float32) / 255.0).unsqueeze(0) |
|
|
|
|
| def resize_image(img: Image.Image, width: int, height: int) -> Image.Image: |
| return img.resize((width, height), resample=Image.LANCZOS) |
|
|
|
|
| def blend_overlay(img_1: Image.Image, img_2: Image.Image) -> Image.Image: |
| arr1 = np.array(img_1).astype(float) / 255.0 |
| arr2 = np.array(img_2).astype(float) / 255.0 |
| mask = arr2 < 0.5 |
| result = np.zeros_like(arr1) |
| result[mask] = 2 * arr1[mask] * arr2[mask] |
| result[~mask] = 1 - 2 * (1 - arr1[~mask]) * (1 - arr2[~mask]) |
| return Image.fromarray(np.clip(result * 255, 0, 255).astype(np.uint8)) |
|
|
|
|
| def fill_mask(width: int, height: int, mask: Image.Image, box=(0, 0), color=0) -> Image.Image: |
| bg = Image.new("L", (width, height), color) |
| bg.paste(mask, box, mask) |
| return bg |
|
|
|
|
| def empty_image(width: int, height: int, batch_size: int = 1) -> torch.Tensor: |
| return torch.zeros([batch_size, height, width, 3]) |
|
|
|
|
| def upscale_mask(mask: torch.Tensor, width: int, height: int) -> torch.Tensor: |
| if mask.ndim == 3: |
| mask = mask.unsqueeze(1) |
| mask = common_upscale(mask, width, height, "bicubic", "disabled") |
| mask = mask.squeeze(1) |
| return mask |
|
|
|
|
| def extract_alpha_mask(image: torch.Tensor) -> torch.Tensor: |
| alpha = image[..., 3] |
| if alpha.max() > 1.0: |
| alpha = alpha / 255.0 |
| if len(alpha.shape) == 4: |
| alpha = alpha[:, :, :, 0] |
| return alpha.unsqueeze(1) if alpha.ndim == 3 else alpha |
|
|
|
|
| def ensure_mask_shape(mask: torch.Tensor | None) -> torch.Tensor | None: |
| if mask is None: |
| return None |
| if mask.ndim == 2: |
| return mask.unsqueeze(0) |
| if mask.ndim == 4 and mask.shape[1] == 1: |
| return mask.squeeze(1) |
| return mask |
|
|
|
|
| COLOR_PRESETS = { |
| "black": "#000000", |
| "white": "#FFFFFF", |
| "red": "#FF0000", |
| "green": "#00FF00", |
| "blue": "#0000FF", |
| "yellow": "#FFFF00", |
| "cyan": "#00FFFF", |
| "magenta": "#FF00FF", |
| "gray": "#808080", |
| "silver": "#C0C0C0", |
| "maroon": "#800000", |
| "olive": "#808000", |
| "purple": "#800080", |
| "teal": "#008080", |
| "navy": "#000080", |
| "orange": "#FFA500", |
| "pink": "#FFC0CB", |
| "brown": "#A52A2A", |
| "violet": "#EE82EE", |
| "indigo": "#4B0082", |
| "light_gray": "#D3D3D3", |
| "dark_gray": "#A9A9A9", |
| "light_blue": "#ADD8E6", |
| "dark_blue": "#00008B", |
| "light_green": "#90EE90", |
| "dark_green": "#006400", |
| } |
|
|
|
|
| def color_format(color: str) -> str: |
| if not color: |
| return "" |
|
|
| color = color.strip().upper() |
| if not color.startswith("#"): |
| color = f"#{color}" |
|
|
| color = color[1:] |
| if len(color) == 3: |
| r, g, b = color[0], color[1], color[2] |
| return f"#{r}{r}{g}{g}{b}{b}" |
| if len(color) < 6: |
| raise ValueError(f"Invalid color format: {color}") |
| if len(color) > 6: |
| color = color[:6] |
|
|
| return f"#{color}" |
|
|
