File size: 3,363 Bytes
c6535db | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | # Shared utility helpers
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}"
|