1499 lines
41 KiB
Python
1499 lines
41 KiB
Python
"""Support utilities for the nodes package."""
|
|
|
|
import base64
|
|
import contextlib
|
|
import functools
|
|
import importlib
|
|
import io
|
|
import itertools
|
|
import json
|
|
import math
|
|
import operator
|
|
import os
|
|
import shlex
|
|
import shutil
|
|
import socket
|
|
import subprocess
|
|
import sys
|
|
import uuid
|
|
from collections.abc import Callable, Sequence
|
|
from enum import Enum
|
|
from functools import reduce
|
|
from pathlib import Path
|
|
from typing import TypeVar
|
|
|
|
import folder_paths
|
|
import numpy as np
|
|
import numpy.typing as npt
|
|
import requests
|
|
import torch
|
|
from PIL import Image
|
|
|
|
from .install import pip_map
|
|
|
|
try:
|
|
from .log import log
|
|
except ImportError:
|
|
try:
|
|
from log import log
|
|
|
|
log.warn("Imported log without relative path")
|
|
except ImportError:
|
|
import logging
|
|
|
|
log = logging.getLogger("comfy mtb utils")
|
|
log.warn("[comfy mtb] You probably called the file outside a module.")
|
|
|
|
try:
|
|
import open3d as o3d
|
|
except:
|
|
log.warn("You do not have open3D installed, 3d utils won't work")
|
|
o3d = {}
|
|
|
|
|
|
# region SANITY_CHECK Utilities
|
|
|
|
|
|
def make_report():
|
|
"""Generate user report about the mtb installation."""
|
|
pass
|
|
|
|
|
|
# endregion
|
|
|
|
|
|
# region NFOV
|
|
class numpy_NFOV:
|
|
def __init__(self, fov=None, height: int = 400, width: int = 800):
|
|
self.field_of_view = fov or [0.45, 0.45]
|
|
self.PI = np.pi
|
|
self.PI_2 = np.pi * 0.5
|
|
self.PI2 = np.pi * 2.0
|
|
self.height = height
|
|
self.width = width
|
|
self.screen_points = self._get_screen_img()
|
|
|
|
def _get_coord_rad(self, is_center_point, center_point=None):
|
|
if is_center_point:
|
|
center_point = np.array(center_point)
|
|
return (center_point * 2 - 1) * np.array([self.PI, self.PI_2])
|
|
else:
|
|
return (
|
|
(self.screen_points * 2 - 1)
|
|
* np.array([self.PI, self.PI_2])
|
|
* (np.ones(self.screen_points.shape) * self.field_of_view)
|
|
)
|
|
|
|
def _get_screen_img(self):
|
|
xx, yy = np.meshgrid(
|
|
np.linspace(0, 1, self.width), np.linspace(0, 1, self.height)
|
|
)
|
|
return np.array([xx.ravel(), yy.ravel()]).T
|
|
|
|
def _calc_spherical_to_gnomonic(self, converted_screen_coord):
|
|
x = converted_screen_coord.T[0]
|
|
y = converted_screen_coord.T[1]
|
|
|
|
rou = np.sqrt(x**2 + y**2)
|
|
c = np.arctan(rou)
|
|
sin_c = np.sin(c)
|
|
cos_c = np.cos(c)
|
|
|
|
lat = np.arcsin(
|
|
cos_c * np.sin(self.cp[1]) + (y * sin_c * np.cos(self.cp[1])) / rou
|
|
)
|
|
lon = self.cp[0] + np.arctan2(
|
|
x * sin_c,
|
|
rou * np.cos(self.cp[1]) * cos_c - y * np.sin(self.cp[1]) * sin_c,
|
|
)
|
|
|
|
lat = (lat / self.PI_2 + 1.0) * 0.5
|
|
lon = (lon / self.PI + 1.0) * 0.5
|
|
|
|
return np.array([lon, lat]).T
|
|
|
|
def _bilinear_interpolation(self, screen_coord):
|
|
uf = np.mod(screen_coord.T[0], 1) * self.frame_width # long - width
|
|
vf = np.mod(screen_coord.T[1], 1) * self.frame_height # lat - height
|
|
|
|
x0 = np.floor(uf).astype(int) # coord of pixel to bottom left
|
|
y0 = np.floor(vf).astype(int)
|
|
x2 = np.add(
|
|
x0, np.ones(uf.shape).astype(int)
|
|
) # coords of pixel to top right
|
|
y2 = np.add(y0, np.ones(vf.shape).astype(int))
|
|
|
|
base_y0 = np.multiply(y0, self.frame_width)
|
|
base_y2 = np.multiply(y2, self.frame_width)
|
|
|
|
A_idx = np.add(base_y0, x0)
|
|
B_idx = np.add(base_y2, x0)
|
|
C_idx = np.add(base_y0, x2)
|
|
D_idx = np.add(base_y2, x2)
|
|
|
|
flat_img = np.reshape(self.frame, [-1, self.frame_channel])
|
|
|
|
A = np.take(flat_img, A_idx, axis=0)
|
|
B = np.take(flat_img, B_idx, axis=0)
|
|
C = np.take(flat_img, C_idx, axis=0)
|
|
D = np.take(flat_img, D_idx, axis=0)
|
|
|
|
wa = np.multiply(x2 - uf, y2 - vf)
|
|
wb = np.multiply(x2 - uf, vf - y0)
|
|
wc = np.multiply(uf - x0, y2 - vf)
|
|
wd = np.multiply(uf - x0, vf - y0)
|
|
|
|
# interpolate
|
|
AA = np.multiply(A, np.array([wa, wa, wa]).T)
|
|
BB = np.multiply(B, np.array([wb, wb, wb]).T)
|
|
CC = np.multiply(C, np.array([wc, wc, wc]).T)
|
|
DD = np.multiply(D, np.array([wd, wd, wd]).T)
|
|
nfov = np.reshape(
|
|
np.round(AA + BB + CC + DD).astype(np.uint8),
|
|
[self.height, self.width, 3],
|
|
)
|
|
|
|
return nfov
|
|
|
|
def to_nfov(self, frame, center_point):
|
|
self.frame = frame
|
|
self.frame_height = frame.shape[0]
|
|
self.frame_width = frame.shape[1]
|
|
self.frame_channel = frame.shape[2]
|
|
|
|
self.cp = self._get_coord_rad(
|
|
center_point=center_point, is_center_point=True
|
|
)
|
|
converted_screen_coord = self._get_coord_rad(is_center_point=False)
|
|
return self._bilinear_interpolation(
|
|
self._calc_spherical_to_gnomonic(converted_screen_coord)
|
|
)
|
|
|
|
|
|
# endregion
|
|
|
|
|
|
# region SERVER Utilities
|
|
class IPChecker:
|
|
"""Checks for a working ip address used by the server."""
|
|
|
|
def __init__(self):
|
|
self.ips = list(self.get_local_ips())
|
|
log.debug(f"Found {len(self.ips)} local ips")
|
|
self.checked_ips: set[str] = set()
|
|
|
|
def get_working_ip(self, test_url_template: str):
|
|
for ip in self.ips:
|
|
if ip not in self.checked_ips:
|
|
self.checked_ips.add(ip)
|
|
test_url = test_url_template.format(ip)
|
|
if self._test_url(test_url):
|
|
return ip
|
|
return None
|
|
|
|
@staticmethod
|
|
def get_local_ips(prefix: str = "192.168."):
|
|
hostname = socket.gethostname()
|
|
log.debug(f"Getting local ips for {hostname}")
|
|
for info in socket.getaddrinfo(hostname, None):
|
|
# Filter out IPv6 addresses if you only want IPv4
|
|
log.debug(info)
|
|
# if info[1] == socket.SOCK_STREAM and
|
|
if info[0] == socket.AF_INET and info[4][0].startswith(prefix):
|
|
yield info[4][0]
|
|
|
|
def _test_url(self, url: str):
|
|
try:
|
|
response = requests.get(url, timeout=10)
|
|
return response.status_code == 200
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
@functools.lru_cache(maxsize=1)
|
|
def get_server_info():
|
|
"""Get server info."""
|
|
from comfy.cli_args import args
|
|
|
|
ip_checker = IPChecker()
|
|
base_url: str = args.listen
|
|
if base_url == "0.0.0.0": # noqa: S104
|
|
log.debug("Server set to 0.0.0.0, we will try to resolve the host IP")
|
|
base_url = ip_checker.get_working_ip(
|
|
f"http://{{}}:{args.port}/history"
|
|
)
|
|
log.debug(f"Setting ip to {base_url}")
|
|
return (base_url, args.port)
|
|
|
|
|
|
# endregion
|
|
|
|
|
|
# region MISC Utilities
|
|
def glob_multiple(
|
|
path: Path, patterns: list[str], recursive: bool = False
|
|
) -> list[Path]:
|
|
"""Combine multiple glob patterns into a single iterator."""
|
|
return list(reduce(operator.or_, (set(path.glob(p)) for p in patterns)))
|
|
|
|
|
|
def build_glob_patterns(
|
|
extensions: list[str], recursive: bool = False
|
|
) -> list[str]:
|
|
"""Build glob patterns for given extensions."""
|
|
prefix = "**/" if recursive else ""
|
|
return [f"{prefix}*.{ext}" for ext in extensions]
|
|
|
|
|
|
class SortMode(Enum):
|
|
NONE = "none"
|
|
MODIFIED = "modified"
|
|
MODIFIED_REVERSE = "modified-reverse"
|
|
NAME = "name"
|
|
NAME_REVERSE = "name-reverse"
|
|
|
|
@classmethod
|
|
def from_str(cls, value: str | None) -> "SortMode|None":
|
|
if not value:
|
|
return None
|
|
try:
|
|
return cls(value.lower())
|
|
except ValueError:
|
|
log.warning(f"Sort mode {value} not supported")
|
|
return None
|
|
|
|
|
|
# TODO: use mtb.core directly instead of copying parts here
|
|
T = TypeVar("T", bound="StringConvertibleEnum")
|
|
|
|
|
|
class StringConvertibleEnum(Enum):
|
|
"""Base class for enums with utility methods for string conversion and member listing."""
|
|
|
|
@classmethod
|
|
def from_str(cls: type[T], label: str | T) -> T:
|
|
"""
|
|
Convert a string to the corresponding enum value (case sensitive).
|
|
|
|
Args:
|
|
label (Union[str, T]): The string or enum value to convert.
|
|
|
|
Returns
|
|
-------
|
|
T: The corresponding enum value.
|
|
|
|
Raises
|
|
------
|
|
ValueError: If the label does not correspond to any enum member.
|
|
"""
|
|
if isinstance(label, cls):
|
|
return label
|
|
if isinstance(label, str):
|
|
# from key
|
|
if label in cls.__members__:
|
|
return cls[label]
|
|
|
|
for member in cls:
|
|
if member.value == label:
|
|
return member
|
|
|
|
raise ValueError(
|
|
f"Unknown label: '{label}'. Valid members: {list(cls.__members__.keys())}, "
|
|
f"valid values: {cls.list_members()}"
|
|
)
|
|
|
|
@classmethod
|
|
def to_str(cls: type[T], enum_value: T) -> str:
|
|
"""
|
|
Convert an enum value to its string representation.
|
|
|
|
Args:
|
|
enum_value (T): The enum value to convert.
|
|
|
|
Returns
|
|
-------
|
|
str: The string representation of the enum value.
|
|
|
|
Raises
|
|
------
|
|
ValueError: If the enum value is invalid.
|
|
"""
|
|
if isinstance(enum_value, cls):
|
|
return enum_value.value
|
|
raise ValueError(f"Invalid Enum: {enum_value}")
|
|
|
|
@classmethod
|
|
def list_members(cls: type[T]) -> list[str]:
|
|
"""
|
|
Return a list of string representations of all enum members.
|
|
|
|
Returns
|
|
-------
|
|
List[str]: List of all enum member values.
|
|
"""
|
|
return [enum.value for enum in cls]
|
|
|
|
def __str__(self) -> str:
|
|
"""
|
|
Returns the string representation of the enum value.
|
|
|
|
Returns
|
|
-------
|
|
str: The string representation of the enum value.
|
|
"""
|
|
return self.value
|
|
|
|
|
|
class Precision(StringConvertibleEnum):
|
|
FULL = "full"
|
|
FP32 = "fp32"
|
|
FP16 = "fp16"
|
|
BF16 = "bf16"
|
|
FP8 = "fp8"
|
|
|
|
def to_dtype(self):
|
|
match self:
|
|
case Precision.FP32 | Precision.FULL:
|
|
return torch.float32
|
|
case Precision.FP16:
|
|
return torch.float16
|
|
case Precision.BF16:
|
|
return torch.bfloat16
|
|
case Precision.FP8:
|
|
return torch.float8_e4m3fn
|
|
|
|
|
|
class Operation(StringConvertibleEnum):
|
|
COPY = "copy"
|
|
CONVERT = "convert"
|
|
DELETE = "delete"
|
|
|
|
|
|
def backup_file(
|
|
fp: Path,
|
|
target: Path | None = None,
|
|
backup_dir: str = ".bak",
|
|
suffix: str | None = None,
|
|
prefix: str | None = None,
|
|
):
|
|
"""Backup a file by copying it to a backup directory."""
|
|
if not fp.exists():
|
|
raise FileNotFoundError(f"No file found at {fp}")
|
|
|
|
backup_directory = target or fp.parent / backup_dir
|
|
backup_directory.mkdir(parents=True, exist_ok=True)
|
|
|
|
stem = fp.stem
|
|
|
|
if suffix or prefix:
|
|
new_stem = f"{prefix or ''}{stem}{suffix or ''}"
|
|
else:
|
|
new_stem = f"{stem}_{uuid.uuid4()}"
|
|
|
|
backup_file_path = backup_directory / f"{new_stem}{fp.suffix}"
|
|
|
|
# Perform the backup
|
|
shutil.copy(fp, backup_file_path)
|
|
log.debug(f"File backed up to {backup_file_path}")
|
|
|
|
|
|
def hex_to_rgb(hex_color, *, bgr=False):
|
|
"""Convert hex to rgb tuple."""
|
|
try:
|
|
hex_color = hex_color.lstrip("#")
|
|
if bgr:
|
|
return tuple(int(hex_color[i : i + 2], 16) for i in (4, 2, 0))
|
|
return tuple(int(hex_color[i : i + 2], 16) for i in (0, 2, 4))
|
|
|
|
except ValueError:
|
|
log.error(f"Invalid hex color: {hex_color}")
|
|
return (0, 0, 0)
|
|
|
|
|
|
def add_path(path: str | Path | List[str] | List[Path], *, prepend=False):
|
|
"""Add a path(s) to the system path.
|
|
|
|
Supports Path, str, and list (of either).
|
|
"""
|
|
if isinstance(path, list):
|
|
for p in path:
|
|
add_path(p, prepend=prepend)
|
|
return
|
|
|
|
if isinstance(path, Path):
|
|
path = path.resolve().as_posix()
|
|
|
|
if path not in sys.path:
|
|
if prepend:
|
|
sys.path.insert(0, path)
|
|
else:
|
|
sys.path.append(path)
|
|
|
|
|
|
def run_command(cmd, ignored_lines_start: Optional[str | List[str]] = None):
|
|
"""Unsafe command runner."""
|
|
if ignored_lines_start is None:
|
|
ignored_lines_start = []
|
|
|
|
if isinstance(cmd, str):
|
|
shell_cmd = cmd
|
|
elif isinstance(cmd, list):
|
|
shell_cmd = " ".join(
|
|
arg.as_posix() if isinstance(arg, Path) else shlex.quote(str(arg))
|
|
for arg in cmd
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
f"Invalid 'cmd' argument. \
|
|
It must be a string or a list of arguments. {type(cmd)}"
|
|
)
|
|
|
|
try:
|
|
_run_command(shell_cmd, ignored_lines_start)
|
|
except subprocess.CalledProcessError as e:
|
|
print(
|
|
f"Command failed with return code: {e.returncode}", file=sys.stderr
|
|
)
|
|
print(e.stderr.strip(), file=sys.stderr)
|
|
|
|
except KeyboardInterrupt:
|
|
print("Command execution interrupted.")
|
|
|
|
|
|
def _run_command(shell_cmd, ignored_lines_start):
|
|
log.debug(f"Running {shell_cmd}")
|
|
|
|
result = subprocess.run(
|
|
shell_cmd,
|
|
capture_output=True,
|
|
text=True,
|
|
shell=True, # noqa: S602
|
|
check=True,
|
|
)
|
|
|
|
stdout_lines = result.stdout.strip().split("\n")
|
|
stderr_lines = result.stderr.strip().split("\n")
|
|
|
|
# Print stdout, skipping ignored lines
|
|
for line in stdout_lines:
|
|
if not any(line.startswith(ign) for ign in ignored_lines_start):
|
|
print(line)
|
|
|
|
# Print stderr
|
|
for line in stderr_lines:
|
|
print(line, file=sys.stderr)
|
|
|
|
print("Command executed successfully!")
|
|
|
|
|
|
def import_install(package_name):
|
|
"""Install and import a package if it is not already imported."""
|
|
package_spec = reqs_map.get(package_name, package_name)
|
|
|
|
try:
|
|
importlib.import_module(package_name)
|
|
|
|
except Exception: # (ImportError, ModuleNotFoundError):
|
|
run_command(
|
|
[
|
|
Path(sys.executable).as_posix(),
|
|
"-m",
|
|
"pip",
|
|
"install",
|
|
package_spec,
|
|
]
|
|
)
|
|
importlib.import_module(package_name)
|
|
|
|
|
|
# endregion
|
|
|
|
|
|
# region GLOBAL VARIABLES
|
|
# - detect mode
|
|
comfy_mode = None
|
|
if os.environ.get("COLAB_GPU"):
|
|
comfy_mode = "colab"
|
|
elif "python_embeded" in sys.executable:
|
|
comfy_mode = "embeded"
|
|
elif ".venv" in sys.executable:
|
|
comfy_mode = "venv"
|
|
|
|
# - Get the absolute path of the parent directory of the current script
|
|
here = Path(__file__).parent.absolute()
|
|
|
|
# - Construct the absolute path to the ComfyUI directory
|
|
comfy_dir = Path(folder_paths.base_path)
|
|
models_dir = Path(folder_paths.models_dir)
|
|
output_dir = Path(folder_paths.output_directory)
|
|
input_dir = Path(folder_paths.input_directory)
|
|
styles_dir = comfy_dir / "styles"
|
|
session_id = str(uuid.uuid4())
|
|
# - Construct the path to the font file
|
|
font_path = here / "data" / "font.ttf"
|
|
|
|
# - Add extern folder to path
|
|
extern_root = here / "extern"
|
|
add_path(extern_root)
|
|
|
|
for pth in extern_root.iterdir():
|
|
if pth.is_dir():
|
|
add_path(pth)
|
|
|
|
# - Add the ComfyUI directory and custom nodes path to the sys.path list
|
|
add_path(comfy_dir)
|
|
add_path(comfy_dir / "custom_nodes")
|
|
|
|
# TODO: use the requirements library
|
|
reqs_map = {value: key for key, value in pip_map.items()}
|
|
|
|
# NOTE: store already logged warnings to only alert once.
|
|
warned_messages: set[str] = set()
|
|
|
|
|
|
PIL_FILTER_MAP = {
|
|
"nearest": Image.Resampling.NEAREST,
|
|
"box": Image.Resampling.BOX,
|
|
"bilinear": Image.Resampling.BILINEAR,
|
|
"hamming": Image.Resampling.HAMMING,
|
|
"bicubic": Image.Resampling.BICUBIC,
|
|
"lanczos": Image.Resampling.LANCZOS,
|
|
}
|
|
# endregion
|
|
|
|
|
|
# region TENSOR Utilities
|
|
def to_numpy(image: torch.Tensor) -> npt.NDArray[np.uint8]:
|
|
"""Converts a tensor to a ndarray with proper scaling and type conversion."""
|
|
log.debug(f"Converting tensor to numpy array with shape {image.shape}")
|
|
np_array = np.clip(255.0 * image.cpu().numpy(), 0, 255).astype(np.uint8)
|
|
log.debug(f"Numpy array shape after conversion: {np_array.shape}")
|
|
return np_array
|
|
|
|
|
|
def handle_batch(
|
|
tensor: torch.Tensor,
|
|
func: Callable[[torch.Tensor], Image.Image | npt.NDArray[np.uint8]],
|
|
) -> list[Image.Image] | list[npt.NDArray[np.uint8]]:
|
|
"""Handles batch processing for a given tensor and conversion function."""
|
|
return [func(tensor[i]) for i in range(tensor.shape[0])]
|
|
|
|
|
|
def tensor2pil(tensor: torch.Tensor) -> list[Image.Image]:
|
|
"""Converts a batch of tensors to a list of PIL Images."""
|
|
|
|
def single_tensor2pil(t: torch.Tensor) -> Image.Image:
|
|
np_array = to_numpy(t)
|
|
if np_array.ndim == 2: # (H, W) for masks
|
|
return Image.fromarray(np_array, mode="L")
|
|
elif np_array.ndim == 3: # (H, W, C) for RGB/RGBA
|
|
if np_array.shape[2] == 3:
|
|
return Image.fromarray(np_array, mode="RGB")
|
|
elif np_array.shape[2] == 4:
|
|
return Image.fromarray(np_array, mode="RGBA")
|
|
raise ValueError(f"Invalid tensor shape: {t.shape}")
|
|
|
|
return handle_batch(tensor, single_tensor2pil)
|
|
|
|
|
|
def pil2tensor(images: Image.Image | list[Image.Image]) -> torch.Tensor:
|
|
"""Converts a PIL Image or a list of PIL Images to a tensor."""
|
|
|
|
def single_pil2tensor(image: Image.Image) -> torch.Tensor:
|
|
np_image = np.array(image).astype(np.float32) / 255.0
|
|
if np_image.ndim == 2: # Grayscale
|
|
return torch.from_numpy(np_image).unsqueeze(0) # (1, H, W)
|
|
else: # RGB or RGBA
|
|
return torch.from_numpy(np_image).unsqueeze(0) # (1, H, W, C)
|
|
|
|
if isinstance(images, Image.Image):
|
|
return single_pil2tensor(images)
|
|
else:
|
|
return torch.cat([single_pil2tensor(img) for img in images], dim=0)
|
|
|
|
|
|
def np2tensor(
|
|
np_array: npt.NDArray[np.float32] | Sequence[npt.NDArray[np.float32]],
|
|
) -> torch.Tensor:
|
|
"""Converts a NumPy array or a list of NumPy arrays to a tensor."""
|
|
|
|
def single_np2tensor(array: npt.NDArray[np.float32]) -> torch.Tensor:
|
|
if array.ndim == 2: # (H, W) for masks
|
|
return torch.from_numpy(
|
|
array.astype(np.float32) / 255.0
|
|
).unsqueeze(0) # (1, H, W)
|
|
elif array.ndim == 3: # (H, W, C) for RGB/RGBA
|
|
return torch.from_numpy(
|
|
array.astype(np.float32) / 255.0
|
|
).unsqueeze(0) # (1, H, W, C)
|
|
raise ValueError(f"Invalid array shape: {array.shape}")
|
|
|
|
if isinstance(np_array, np.ndarray):
|
|
return single_np2tensor(np_array)
|
|
else:
|
|
return torch.cat([single_np2tensor(arr) for arr in np_array], dim=0)
|
|
|
|
|
|
def tensor2np(tensor: torch.Tensor) -> list[npt.NDArray[np.uint8]]:
|
|
"""Converts a batch of tensors to a list of NumPy arrays."""
|
|
|
|
def single_tensor2np(t: torch.Tensor) -> npt.NDArray[np.uint8]:
|
|
t = t.squeeze() # Remove any singleton dimensions
|
|
if t.ndim == 2: # (H, W) for masks
|
|
return to_numpy(t)
|
|
elif t.ndim == 3: # (C, H, W) for RGB/RGBA
|
|
if t.shape[0] in [1, 3, 4]: # Channel-first format
|
|
t = t.permute(1, 2, 0)
|
|
return to_numpy(t)
|
|
else:
|
|
raise ValueError(f"Invalid tensor shape: {t.shape}")
|
|
|
|
return handle_batch(tensor, single_tensor2np)
|
|
|
|
|
|
def nextAvailable(path: Path | str) -> Path:
|
|
"""
|
|
Find the next available path by adding a numbered suffix. (mimics comfy's version).
|
|
|
|
Args:
|
|
path (Path): The original path to check
|
|
|
|
Returns
|
|
-------
|
|
Path: A path that doesn't exist yet
|
|
"""
|
|
path = Path(path)
|
|
|
|
if not path.is_absolute():
|
|
path = output_dir / path
|
|
|
|
if not path.exists():
|
|
return path
|
|
|
|
stem = path.stem
|
|
suffix = path.suffix
|
|
parent = path.parent
|
|
|
|
counter = 1
|
|
while True:
|
|
new_path = parent / f"{stem}_{counter:04d}{suffix}"
|
|
if not new_path.exists():
|
|
return new_path
|
|
counter += 1
|
|
|
|
|
|
def pad(img: np.ndarray, left, right, top, bottom):
|
|
"""Pad an image with wrap padding on edges."""
|
|
pad_width = np.array(((0, 0), (top, bottom), (left, right)))
|
|
print(
|
|
f"pad_width: {pad_width}, shape: {pad_width.shape}"
|
|
) # Debugging line
|
|
return np.pad(img, pad_width, mode="wrap") # type: ignore
|
|
|
|
|
|
def tiles_infer(tiles, ort_session, progress_callback=None):
|
|
"""Infer each tile with the given model.
|
|
|
|
progress_callback will be called with
|
|
arguments : current tile idx and total tiles amount
|
|
(used to show progress on cursor in Blender).
|
|
"""
|
|
out_channels = 3 # normal map RGB channels
|
|
tiles_nb = tiles.shape[0]
|
|
pred_tiles = np.empty(
|
|
(tiles_nb, out_channels, tiles.shape[2], tiles.shape[3])
|
|
)
|
|
pred_tiles = np.empty(
|
|
(tiles_nb, out_channels, tiles.shape[2], tiles.shape[3])
|
|
)
|
|
|
|
for i in range(tiles_nb):
|
|
if progress_callback is not None:
|
|
progress_callback(i + 1, tiles_nb)
|
|
pred_tiles[i] = ort_session.run(
|
|
None, {"input": tiles[i : i + 1].astype(np.float32)}
|
|
)[0]
|
|
|
|
return pred_tiles
|
|
|
|
|
|
def generate_mask(tile_size, stride_size):
|
|
"""Generate a pyramidal-like mask.
|
|
|
|
Used for mixing overlapping predicted tiles.
|
|
"""
|
|
tile_h, tile_w = tile_size
|
|
stride_h, stride_w = stride_size
|
|
ramp_h = tile_h - stride_h
|
|
ramp_w = tile_w - stride_w
|
|
|
|
mask = np.ones((tile_h, tile_w))
|
|
|
|
# ramps in width direction
|
|
mask[ramp_h:-ramp_h, :ramp_w] = np.linspace(0, 1, num=ramp_w)
|
|
mask[ramp_h:-ramp_h, -ramp_w:] = np.linspace(1, 0, num=ramp_w)
|
|
# ramps in height direction
|
|
mask[:ramp_h, ramp_w:-ramp_w] = np.transpose(
|
|
np.linspace(0, 1, num=ramp_h)[None],
|
|
(1, 0), # type: ignore
|
|
)
|
|
mask[-ramp_h:, ramp_w:-ramp_w] = np.transpose(
|
|
np.linspace(1, 0, num=ramp_h)[None],
|
|
(1, 0), # type: ignore
|
|
)
|
|
|
|
# Assume tiles are squared
|
|
if ramp_h != ramp_w:
|
|
raise ValueError("Ramps in both height and width directions required")
|
|
|
|
# top left corner
|
|
corner = np.rot90(corner_mask(ramp_h), 2) # type: ignore
|
|
mask[:ramp_h, :ramp_w] = corner
|
|
# top right corner
|
|
corner = np.flip(corner, 1) # type: ignore
|
|
mask[:ramp_h, -ramp_w:] = corner
|
|
# bottom right corner
|
|
corner = np.flip(corner, 0) # type: ignore
|
|
mask[-ramp_h:, -ramp_w:] = corner
|
|
# bottom right corner
|
|
corner = np.flip(corner, 1) # type: ignore
|
|
mask[-ramp_h:, :ramp_w] = corner
|
|
|
|
return mask
|
|
|
|
|
|
def corner_mask(side_length):
|
|
"""Generate the corner part of the pyramidal-like mask.
|
|
|
|
Currently, only for square shapes.
|
|
"""
|
|
corner = np.zeros([side_length, side_length])
|
|
|
|
for h in range(0, side_length):
|
|
for w in range(0, side_length):
|
|
if h >= w:
|
|
sh = h / (side_length - 1)
|
|
corner[h, w] = 1 - sh
|
|
if h <= w:
|
|
sw = w / (side_length - 1)
|
|
corner[h, w] = 1 - sw
|
|
|
|
return corner - 0.25 * scaling_mask(side_length)
|
|
|
|
|
|
def scaling_mask(side_length: int):
|
|
"""Generate a scaling mask for image blending."""
|
|
scaling = np.zeros([side_length, side_length])
|
|
|
|
for h in range(0, side_length):
|
|
for w in range(0, side_length):
|
|
sh = h / (side_length - 1)
|
|
sw = w / (side_length - 1)
|
|
if h >= w and h <= side_length - w:
|
|
scaling[h, w] = sw
|
|
if h <= w and h <= side_length - w:
|
|
scaling[h, w] = sh
|
|
if h >= w and h >= side_length - w:
|
|
scaling[h, w] = 1 - sh
|
|
if h <= w and h >= side_length - w:
|
|
scaling[h, w] = 1 - sw
|
|
|
|
return 2 * scaling
|
|
|
|
|
|
def _validate_stride(w, h, tile_w, tile_h):
|
|
if (h % 2 != 0) or (w % 2 != 0):
|
|
raise ValueError("Stride must be even")
|
|
|
|
if (h < tile_h / 2) or (w < tile_w / 2):
|
|
raise ValueError("Stride must be greater or equal than half tile size")
|
|
|
|
if (h > tile_h) or (w > tile_w):
|
|
raise ValueError("Stride must be smaller or equal than tile size")
|
|
|
|
|
|
def tiles_merge(tiles, stride_size, img_size, paddings):
|
|
"""Merge the list of tiles into one image.
|
|
|
|
img_size is the original size, before padding.
|
|
"""
|
|
_, tile_h, tile_w = tiles[0].shape
|
|
pad_left, pad_right, pad_top, pad_bottom = paddings
|
|
height = img_size[1] + pad_top + pad_bottom
|
|
width = img_size[2] + pad_left + pad_right
|
|
stride_h, stride_w = stride_size
|
|
|
|
_validate_stride(stride_w, stride_h, tile_w, tile_h)
|
|
|
|
merged = np.zeros((img_size[0], height, width))
|
|
mask = generate_mask((tile_h, tile_w), stride_size)
|
|
|
|
h_range = ((height - tile_h) // stride_h) + 1
|
|
w_range = ((width - tile_w) // stride_w) + 1
|
|
|
|
idx = 0
|
|
for h in range(0, h_range):
|
|
for w in range(0, w_range):
|
|
h_from, h_to = h * stride_h, h * stride_h + tile_h
|
|
w_from, w_to = w * stride_w, w * stride_w + tile_w
|
|
merged[:, h_from:h_to, w_from:w_to] += tiles[idx] * mask
|
|
idx += 1
|
|
|
|
return merged[:, pad_top:-pad_bottom, pad_left:-pad_right]
|
|
|
|
|
|
def tiles_split(img, tile_size, stride_size):
|
|
"""Return list of tiles from the given image.
|
|
|
|
Also returns the padding used to fit the tiles in it.
|
|
Input image must have dimension C,H,W.
|
|
"""
|
|
log.debug(f"Splitting img: tile {tile_size}, stride {stride_size} ")
|
|
tile_h, tile_w = tile_size
|
|
stride_h, stride_w = stride_size
|
|
img_h, img_w = img.shape[0], img.shape[1]
|
|
|
|
_validate_stride(stride_w, stride_h, tile_w, tile_h)
|
|
|
|
# find total height & width padding sizes
|
|
pad_h, pad_w = 0, 0
|
|
remainer_h = (img_h - tile_h) % stride_h
|
|
remainer_w = (img_w - tile_w) % stride_w
|
|
if remainer_h != 0:
|
|
pad_h = stride_h - remainer_h
|
|
if remainer_w != 0:
|
|
pad_w = stride_w - remainer_w
|
|
|
|
# if tile bigger than image, pad image to tile size
|
|
if tile_h > img_h:
|
|
pad_h = tile_h - img_h
|
|
if tile_w > img_w:
|
|
pad_w = tile_w - img_w
|
|
|
|
# pad image, add extra stride to padding to avoid pyramid
|
|
# weighting leaking onto the valid part of the picture
|
|
pad_left = pad_w // 2 + stride_w
|
|
pad_right = pad_left if pad_w % 2 == 0 else pad_left + 1
|
|
pad_top = pad_h // 2 + stride_h
|
|
pad_bottom = pad_top if pad_h % 2 == 0 else pad_top + 1
|
|
img = pad(img, pad_left, pad_right, pad_top, pad_bottom)
|
|
img_h, img_w = img.shape[1], img.shape[2]
|
|
|
|
# extract tiles
|
|
h_range = ((img_h - tile_h) // stride_h) + 1
|
|
w_range = ((img_w - tile_w) // stride_w) + 1
|
|
tiles = np.empty([h_range * w_range, img.shape[0], tile_h, tile_w])
|
|
idx = 0
|
|
for h in range(0, h_range):
|
|
for w in range(0, w_range):
|
|
h_from, h_to = h * stride_h, h * stride_h + tile_h
|
|
w_from, w_to = w * stride_w, w * stride_w + tile_w
|
|
tiles[idx] = img[:, h_from:h_to, w_from:w_to]
|
|
idx += 1
|
|
|
|
return tiles, (pad_left, pad_right, pad_top, pad_bottom)
|
|
|
|
|
|
# endregion
|
|
|
|
|
|
# region MODEL Utilities
|
|
def download_antelopev2():
|
|
"""Dumb hardcoded method to download antelopev2 model.
|
|
|
|
TODO: remove and use the system from download_models.py
|
|
"""
|
|
antelopev2_url = (
|
|
"https://drive.google.com/uc?id=18wEUfMNohBJ4K3Ly5wpTejPfDzp-8fI8"
|
|
)
|
|
|
|
try:
|
|
import gdown
|
|
|
|
log.debug("Loading antelopev2 model")
|
|
|
|
dest = get_model_path("insightface")
|
|
archive = dest / "antelopev2.zip"
|
|
final_path = dest / "models" / "antelopev2"
|
|
if not final_path.exists():
|
|
log.info(f"antelopev2 not found, downloading to {dest}")
|
|
gdown.download(
|
|
antelopev2_url,
|
|
archive.as_posix(),
|
|
resume=True,
|
|
)
|
|
|
|
log.info(f"Unzipping antelopev2 to {final_path}")
|
|
|
|
if archive.exists():
|
|
# we unzip it
|
|
import zipfile
|
|
|
|
with zipfile.ZipFile(archive.as_posix(), "r") as zip_ref:
|
|
zip_ref.extractall(final_path.parent.as_posix())
|
|
|
|
except Exception as e:
|
|
log.error(
|
|
f"Could not load or download antelopev2 model, \
|
|
download it manually from {antelopev2_url}"
|
|
)
|
|
raise e
|
|
|
|
|
|
def get_model_path(fam, model=None):
|
|
"""Get any model path given it's family and name.
|
|
|
|
This supports extra_models and models folders.
|
|
"""
|
|
log.debug(f"Requesting {fam} with model {model}")
|
|
res = None
|
|
if model:
|
|
res = folder_paths.get_full_path(fam, model)
|
|
else:
|
|
# this one can raise errors...
|
|
with contextlib.suppress(KeyError):
|
|
res = folder_paths.get_folder_paths(fam)
|
|
|
|
if res:
|
|
if isinstance(res, list):
|
|
if len(res) > 1:
|
|
warn_msg = f"Found multiple match, we will pick the last {res[-1]}\n{res}"
|
|
if warn_msg not in warned_messages:
|
|
log.info(warn_msg)
|
|
warned_messages.add(warn_msg)
|
|
res = res[-1]
|
|
res = Path(res)
|
|
log.debug(f"Resolved model path from folder_paths: {res}")
|
|
else:
|
|
res = models_dir / fam
|
|
if model:
|
|
res /= model
|
|
|
|
return res
|
|
|
|
|
|
# endregion
|
|
|
|
# region geo
|
|
|
|
|
|
def euler_to_rotation_matrix(
|
|
x_deg: float, y_deg: float, z_deg: float
|
|
) -> np.ndarray[np.float64]:
|
|
# Convert degrees to radians
|
|
x = np.radians(x_deg)
|
|
y = np.radians(y_deg)
|
|
z = np.radians(z_deg)
|
|
|
|
# Rotation matrix around x-axis
|
|
Rx = np.array(
|
|
[[1, 0, 0], [0, np.cos(x), -np.sin(x)], [0, np.sin(x), np.cos(x)]]
|
|
)
|
|
|
|
# Rotation matrix around y-axis
|
|
Ry = np.array(
|
|
[[np.cos(y), 0, np.sin(y)], [0, 1, 0], [-np.sin(y), 0, np.cos(y)]]
|
|
)
|
|
|
|
# Rotation matrix around z-axis
|
|
Rz = np.array(
|
|
[[np.cos(z), -np.sin(z), 0], [np.sin(z), np.cos(z), 0], [0, 0, 1]]
|
|
)
|
|
|
|
return Rz @ Ry @ Rx
|
|
|
|
|
|
def rotate_mesh(mesh, x_deg: float, y_deg: float, z_deg: float, center=None):
|
|
R = euler_to_rotation_matrix(x_deg, y_deg, z_deg)
|
|
return mesh.rotate(R, center) if center is not None else mesh.rotate(R)
|
|
|
|
|
|
def get_transformation_matrix(position, rotation, scale):
|
|
# translation matrix
|
|
T = np.eye(4)
|
|
T[:3, 3] = position
|
|
|
|
# rotation matrix from Euler angles
|
|
R = euler_to_rotation_matrix(*rotation)
|
|
R_homo = np.eye(4)
|
|
R_homo[:3, :3] = R
|
|
|
|
# scaling matrix
|
|
S = np.eye(4)
|
|
S[0, 0] = scale[0]
|
|
S[1, 1] = scale[1]
|
|
S[2, 2] = scale[2]
|
|
|
|
# combine
|
|
return T @ R_homo @ S
|
|
|
|
|
|
def spread_geo(geo, *, cp=False):
|
|
"""Spreads a GEOMETRY type into (mesh,material)."""
|
|
mesh = geo["mesh"] if not cp else copy.copy(geo["mesh"])
|
|
material = geo.get("material", {})
|
|
return (mesh, material)
|
|
|
|
|
|
def json_to_mesh(json_data: str):
|
|
"""Convert JSON to an Open3D mesh."""
|
|
data = json.loads(json_data)
|
|
mesh = o3d.geometry.TriangleMesh()
|
|
|
|
if "vertices" in data:
|
|
mesh.vertices = o3d.utility.Vector3dVector(
|
|
np.array(data["vertices"]).reshape(-1, 3)
|
|
)
|
|
|
|
if "triangles" in data:
|
|
mesh.triangles = o3d.utility.Vector3iVector(
|
|
np.array(data["triangles"]).reshape(-1, 3)
|
|
)
|
|
|
|
if "vertex_normals" in data:
|
|
mesh.vertex_normals = o3d.utility.Vector3dVector(
|
|
np.array(data["vertex_normals"]).reshape(-1, 3)
|
|
)
|
|
|
|
if "vertex_colors" in data:
|
|
mesh.vertex_colors = o3d.utility.Vector3dVector(
|
|
np.array(data["vertex_colors"]).reshape(-1, 3)
|
|
)
|
|
|
|
if "triangle_uvs" in data:
|
|
mesh.triangle_uvs = o3d.utility.Vector2dVector(
|
|
np.array(data["triangle_uvs"]).reshape(-1, 2)
|
|
)
|
|
|
|
return mesh
|
|
|
|
|
|
def mesh_to_json(mesh: o3d.geometry.MeshBase):
|
|
"""Convert an Open3D mesh to JSON."""
|
|
mesh_dict = {
|
|
"vertices": np.asarray(mesh.vertices).tolist(),
|
|
"triangles": np.asarray(mesh.triangles).tolist(),
|
|
}
|
|
|
|
if mesh.has_vertex_normals():
|
|
mesh_dict["vertex_normals"] = np.asarray(mesh.vertex_normals).tolist()
|
|
|
|
if mesh.has_vertex_colors():
|
|
mesh_dict["vertex_colors"] = np.asarray(mesh.vertex_colors).tolist()
|
|
|
|
if mesh.has_triangle_uvs():
|
|
mesh_dict["triangle_uvs"] = np.asarray(mesh.triangle_uvs).tolist()
|
|
|
|
return json.dumps(mesh_dict)
|
|
|
|
|
|
def create_grid(scale=(1, 1, 1), rows=10, columns=10):
|
|
dx, dy, dz = scale
|
|
|
|
# Create vertices
|
|
vertices = []
|
|
for i in np.linspace(-dy / 2, dy / 2, rows + 1):
|
|
vertices.extend(
|
|
[j, 0, i] for j in np.linspace(-dx / 2, dx / 2, columns + 1)
|
|
)
|
|
# Generate triangles
|
|
triangles = []
|
|
for i, j in itertools.product(range(rows), range(columns)):
|
|
p1 = i * (columns + 1) + j
|
|
p2 = i * (columns + 1) + j + 1
|
|
p3 = (i + 1) * (columns + 1) + j
|
|
p4 = (i + 1) * (columns + 1) + j + 1
|
|
|
|
triangles.extend(([p1, p2, p3], [p2, p4, p3]))
|
|
vertices = o3d.utility.Vector3dVector(np.array(vertices))
|
|
triangles = o3d.utility.Vector3iVector(np.array(triangles))
|
|
mesh = o3d.geometry.TriangleMesh(vertices, triangles)
|
|
|
|
mesh.compute_vertex_normals()
|
|
|
|
return mesh
|
|
|
|
|
|
def create_box(scale=(1, 1, 1), divisions=(1, 1, 1)):
|
|
dx, dy, dz = scale
|
|
div_x, div_y, div_z = divisions
|
|
|
|
vertices = []
|
|
for i in np.linspace(-dx / 2, dx / 2, div_x + 1):
|
|
for j in np.linspace(-dy / 2, dy / 2, div_y + 1):
|
|
vertices.extend(
|
|
[i, j, k] for k in np.linspace(-dz / 2, dz / 2, div_z + 1)
|
|
)
|
|
# Generate triangles for the box faces
|
|
triangles = []
|
|
for x, y in itertools.product(range(div_x), range(div_y)):
|
|
for z in range(div_z):
|
|
# Define base index for this cube
|
|
base = z * (div_x + 1) * (div_y + 1) + y * (div_x + 1) + x
|
|
|
|
# Indices for the 8 vertices of the cube
|
|
v0 = base
|
|
v1 = base + 1
|
|
v2 = base + (div_x + 1) + 1
|
|
v3 = base + (div_x + 1)
|
|
v4 = base + (div_x + 1) * (div_y + 1)
|
|
v5 = v4 + 1
|
|
v6 = v4 + (div_x + 1) + 1
|
|
v7 = v4 + (div_x + 1)
|
|
|
|
triangles.extend(
|
|
(
|
|
[v0, v1, v2],
|
|
[v2, v3, v0],
|
|
[v4, v5, v6],
|
|
[v6, v7, v4],
|
|
[v0, v3, v7],
|
|
[v7, v4, v0],
|
|
[v1, v5, v6],
|
|
[v6, v2, v1],
|
|
[v0, v4, v5],
|
|
[v5, v1, v0],
|
|
[v3, v2, v6],
|
|
[v6, v7, v3],
|
|
)
|
|
)
|
|
vertices = o3d.utility.Vector3dVector(np.array(vertices))
|
|
triangles = o3d.utility.Vector3iVector(np.array(triangles))
|
|
mesh = o3d.geometry.TriangleMesh(vertices, triangles)
|
|
|
|
mesh.compute_vertex_normals()
|
|
|
|
return mesh
|
|
|
|
|
|
def create_sphere(radius=1, columns=10, rows=10):
|
|
# Create vertex positions
|
|
vertices = []
|
|
for i in range(rows + 1):
|
|
lat = i * np.pi / rows
|
|
sin_lat = np.sin(lat)
|
|
cos_lat = np.cos(lat)
|
|
|
|
for j in range(columns + 1):
|
|
lon = j * 2 * np.pi / columns
|
|
sin_lon = np.sin(lon)
|
|
cos_lon = np.cos(lon)
|
|
|
|
x = radius * cos_lon * sin_lat
|
|
y = radius * cos_lat
|
|
z = radius * sin_lon * sin_lat
|
|
vertices.append([x, y, z])
|
|
|
|
# Create triangles
|
|
triangles = []
|
|
for i in range(rows):
|
|
for j in range(columns):
|
|
p1 = i * (columns + 1) + j
|
|
p2 = i * (columns + 1) + j + 1
|
|
p3 = (i + 1) * (columns + 1) + j
|
|
p4 = (i + 1) * (columns + 1) + j + 1
|
|
|
|
triangles.extend(([p1, p2, p3], [p2, p4, p3]))
|
|
vertices = o3d.utility.Vector3dVector(np.array(vertices))
|
|
triangles = o3d.utility.Vector3iVector(np.array(triangles))
|
|
mesh = o3d.geometry.TriangleMesh(vertices, triangles)
|
|
|
|
# Assigning random colors to vertices
|
|
N = len(vertices)
|
|
mesh.vertex_colors = o3d.utility.Vector3dVector(
|
|
np.random.uniform(0, 1, size=(N, 3))
|
|
)
|
|
mesh.compute_vertex_normals()
|
|
|
|
return mesh
|
|
|
|
|
|
def create_torus(torus_radius=1, ring_radius=0.5, rows=10, columns=10):
|
|
vertices = []
|
|
for i in range(rows + 1):
|
|
theta = i * 2 * np.pi / rows
|
|
cos_theta = np.cos(theta)
|
|
sin_theta = np.sin(theta)
|
|
|
|
circle_center = torus_radius + ring_radius * cos_theta
|
|
|
|
for j in range(columns + 1):
|
|
phi = j * 2 * np.pi / columns
|
|
cos_phi = np.cos(phi)
|
|
sin_phi = np.sin(phi)
|
|
|
|
x = circle_center * cos_phi
|
|
y = ring_radius * sin_theta
|
|
z = circle_center * sin_phi
|
|
vertices.append([x, y, z])
|
|
|
|
triangles = []
|
|
for i in range(rows):
|
|
for j in range(columns):
|
|
p1 = i * (columns + 1) + j
|
|
p2 = i * (columns + 1) + j + 1
|
|
p3 = (i + 1) * (columns + 1) + j
|
|
p4 = (i + 1) * (columns + 1) + j + 1
|
|
|
|
triangles.extend(([p1, p2, p3], [p2, p4, p3]))
|
|
vertices = o3d.utility.Vector3dVector(np.array(vertices))
|
|
triangles = o3d.utility.Vector3iVector(np.array(triangles))
|
|
mesh = o3d.geometry.TriangleMesh(vertices, triangles)
|
|
|
|
mesh.compute_vertex_normals()
|
|
|
|
return mesh
|
|
|
|
|
|
# endregion
|
|
|
|
# region UV Utilities
|
|
|
|
|
|
def create_uv_map_tensor(width=512, height=512):
|
|
"""Create a UV map tensor."""
|
|
# Create u coordinates
|
|
u_coords = torch.linspace(0.0, 1.0, steps=width)
|
|
# Create v coordinates
|
|
v_coords = torch.linspace(0.0, 1.0, steps=height)
|
|
|
|
# Generate UV grid
|
|
u_grid, v_grid = torch.meshgrid(u_coords, v_coords)
|
|
|
|
uv_map = torch.zeros(height, width, 3, dtype=torch.float32)
|
|
# Fill U channel with U grid transposed
|
|
uv_map[:, :, 0] = u_grid.t()
|
|
# Fill V channel with V grid transposed
|
|
uv_map[:, :, 1] = v_grid.t()
|
|
|
|
return uv_map.unsqueeze(0)
|
|
|
|
|
|
# endregion
|
|
|
|
|
|
# region ANIMATION Utilities
|
|
EASINGS = [
|
|
"Linear",
|
|
"Sine In",
|
|
"Sine Out",
|
|
"Sine In/Out",
|
|
"Quart In",
|
|
"Quart Out",
|
|
"Quart In/Out",
|
|
"Cubic In",
|
|
"Cubic Out",
|
|
"Cubic In/Out",
|
|
"Circ In",
|
|
"Circ Out",
|
|
"Circ In/Out",
|
|
"Back In",
|
|
"Back Out",
|
|
"Back In/Out",
|
|
"Elastic In",
|
|
"Elastic Out",
|
|
"Elastic In/Out",
|
|
"Bounce In",
|
|
"Bounce Out",
|
|
"Bounce In/Out",
|
|
]
|
|
|
|
|
|
def apply_easing(value, easing_type):
|
|
if easing_type == "Linear":
|
|
return value
|
|
|
|
# Back easing functions
|
|
def easeInBack(t):
|
|
s = 1.70158
|
|
return t * t * ((s + 1) * t - s)
|
|
|
|
|
|
def _out_back(t):
|
|
s = 1.70158
|
|
return ((t - 1) * t * ((s + 1) * t + s)) + 1
|
|
|
|
|
|
def _inout_back(t):
|
|
s = 1.70158 * 1.525
|
|
if t < 0.5:
|
|
return (t * t * (t * (s + 1) - s)) * 2
|
|
return ((t - 2) * t * ((s + 1) * t + s) + 2) * 2
|
|
|
|
|
|
# Elastic easing functions
|
|
def _in_elastic(t):
|
|
if t == 0:
|
|
return 0
|
|
if t == 1:
|
|
return 1
|
|
p = 0.3
|
|
s = p / 4
|
|
return -(
|
|
math.pow(2, 10 * (t - 1)) * math.sin((t - 1 - s) * (2 * math.pi) / p)
|
|
)
|
|
|
|
|
|
def _out_elastic(t):
|
|
if t == 0:
|
|
return 0
|
|
if t == 1:
|
|
return 1
|
|
p = 0.3
|
|
s = p / 4
|
|
return math.pow(2, -10 * t) * math.sin((t - s) * (2 * math.pi) / p) + 1
|
|
|
|
|
|
def _inout_elastic(t):
|
|
if t == 0:
|
|
return 0
|
|
if t == 1:
|
|
return 1
|
|
p = 0.3 * 1.5
|
|
s = p / 4
|
|
t = t * 2
|
|
if t < 1:
|
|
return -0.5 * (
|
|
math.pow(2, 10 * (t - 1))
|
|
* math.sin((t - 1 - s) * (2 * math.pi) / p)
|
|
)
|
|
return (
|
|
0.5
|
|
* math.pow(2, -10 * (t - 1))
|
|
* math.sin((t - 1 - s) * (2 * math.pi) / p)
|
|
+ 1
|
|
)
|
|
|
|
|
|
# Bounce easing functions
|
|
def _in_bounce(t):
|
|
return 1 - _out_bounce(1 - t)
|
|
|
|
|
|
def _out_bounce(t):
|
|
if t < (1 / 2.75):
|
|
return 7.5625 * t * t
|
|
elif t < (2 / 2.75):
|
|
t -= 1.5 / 2.75
|
|
return 7.5625 * t * t + 0.75
|
|
elif t < (2.5 / 2.75):
|
|
t -= 2.25 / 2.75
|
|
return 7.5625 * t * t + 0.9375
|
|
else:
|
|
t -= 2.625 / 2.75
|
|
return 7.5625 * t * t + 0.984375
|
|
|
|
|
|
def _inout_bounce(t):
|
|
if t < 0.5:
|
|
return _in_bounce(t * 2) * 0.5
|
|
return _out_bounce(t * 2 - 1) * 0.5 + 0.5
|
|
|
|
|
|
# Quart easing functions
|
|
def _in_quart(t):
|
|
return t * t * t * t
|
|
|
|
|
|
def _out_quart(t):
|
|
t -= 1
|
|
return -(t**2 * t * t - 1)
|
|
|
|
|
|
def _inout_quart(t):
|
|
t *= 2
|
|
if t < 1:
|
|
return 0.5 * t * t * t * t
|
|
t -= 2
|
|
return -0.5 * (t**2 * t * t - 2)
|
|
|
|
|
|
# Cubic easing functions
|
|
def _in_cubic(t):
|
|
return t * t * t
|
|
|
|
|
|
def _out_cubic(t):
|
|
t -= 1
|
|
return t**2 * t + 1
|
|
|
|
|
|
def _inout_cubic(t):
|
|
t *= 2
|
|
if t < 1:
|
|
return 0.5 * t * t * t
|
|
t -= 2
|
|
return 0.5 * (t**2 * t + 2)
|
|
|
|
|
|
# Circ easing functions
|
|
def _in_circ(t):
|
|
return -(math.sqrt(1 - t * t) - 1)
|
|
|
|
|
|
def _out_circ(t):
|
|
t -= 1
|
|
return math.sqrt(1 - t**2)
|
|
|
|
|
|
def _inout_circ(t):
|
|
t *= 2
|
|
if t < 1:
|
|
return -0.5 * (math.sqrt(1 - t**2) - 1)
|
|
t -= 2
|
|
return 0.5 * (math.sqrt(1 - t**2) + 1)
|
|
|
|
|
|
# Sine easing functions
|
|
def _in_sine(t):
|
|
return -math.cos(t * (math.pi / 2)) + 1
|
|
|
|
|
|
def _out_sine(t):
|
|
return math.sin(t * (math.pi / 2))
|
|
|
|
|
|
def _inout_sine(t):
|
|
return -0.5 * (math.cos(math.pi * t) - 1)
|
|
|
|
|
|
def apply_easing(value: float, easing_type: str):
|
|
"""Apply easing function to value based on easing type."""
|
|
if easing_type == "Linear":
|
|
return value
|
|
|
|
easing_functions = {
|
|
"Sine In": _in_sine,
|
|
"Sine Out": _out_sine,
|
|
"Sine In/Out": _inout_sine,
|
|
"Quart In": _in_quart,
|
|
"Quart Out": _out_quart,
|
|
"Quart In/Out": _inout_quart,
|
|
"Cubic In": _in_cubic,
|
|
"Cubic Out": _out_cubic,
|
|
"Cubic In/Out": _inout_cubic,
|
|
"Circ In": _in_circ,
|
|
"Circ Out": _out_circ,
|
|
"Circ In/Out": _inout_circ,
|
|
"Back In": _in_back,
|
|
"Back Out": _out_back,
|
|
"Back In/Out": _inout_back,
|
|
"Elastic In": _in_elastic,
|
|
"Elastic Out": _out_elastic,
|
|
"Elastic In/Out": _inout_elastic,
|
|
"Bounce In": _in_bounce,
|
|
"Bounce Out": _out_bounce,
|
|
"Bounce In/Out": _inout_bounce,
|
|
}
|
|
|
|
function_ = easing_functions.get(easing_type)
|
|
if function_:
|
|
return function_(value)
|
|
|
|
log.error(f"Unknown easing type: {easing_type}")
|
|
log.error(f"Available easing types: {list(easing_functions.keys())}")
|
|
raise ValueError(f"Unknown easing type: {easing_type}")
|
|
|
|
|
|
# endregion
|