AI Lab
95d5e5c94f
This file contains utility functions for image processing, including conversions between PIL images and tensors, resizing, blending, and color formatting.
118 lines
3.3 KiB
Python
118 lines
3.3 KiB
Python
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}"
|