215 lines
6.8 KiB
Python
215 lines
6.8 KiB
Python
"""
|
||
Shared rotation utilities for the rotate-image ComfyUI custom nodes.
|
||
|
||
Provides PIL-based image rotation with multiple interpolation algorithms,
|
||
and conversion helpers between ComfyUI tensor format and PIL Images.
|
||
|
||
Notes on PIL rotate() resampling support
|
||
-----------------------------------------
|
||
PIL's Image.rotate() only accepts NEAREST, BILINEAR, and BICUBIC as
|
||
resampling filters. LANCZOS (a.k.a. ANTIALIAS) is supported for
|
||
resize/thumbnail operations but *not* for affine/rotate transforms.
|
||
|
||
To give users the "LANCZOS" quality option we map it to BICUBIC for the
|
||
rotation step (BICUBIC is the highest quality filter that rotate() accepts)
|
||
and document this behaviour clearly.
|
||
"""
|
||
|
||
import math
|
||
from typing import Tuple
|
||
|
||
import numpy as np
|
||
import torch
|
||
from PIL import Image
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Interpolation method registry
|
||
# ---------------------------------------------------------------------------
|
||
|
||
# All four options are exposed to the user in the ComfyUI UI.
|
||
INTERPOLATION_METHODS = {
|
||
"NEAREST": Image.Resampling.NEAREST,
|
||
"BILINEAR": Image.Resampling.BILINEAR,
|
||
"BICUBIC": Image.Resampling.BICUBIC,
|
||
"LANCZOS": Image.Resampling.LANCZOS,
|
||
}
|
||
|
||
# PIL Image.rotate() only supports these three filters.
|
||
# LANCZOS is silently promoted to BICUBIC for the rotation step.
|
||
_ROTATE_SAFE_FILTERS = {
|
||
Image.Resampling.NEAREST,
|
||
Image.Resampling.BILINEAR,
|
||
Image.Resampling.BICUBIC,
|
||
}
|
||
|
||
DEFAULT_INTERPOLATION = "BICUBIC"
|
||
|
||
|
||
def _rotate_resample(resample: Image.Resampling) -> Image.Resampling:
|
||
"""
|
||
Return a resampling filter that PIL's rotate() accepts.
|
||
|
||
LANCZOS is not supported by rotate(); fall back to BICUBIC which
|
||
provides equivalent quality for affine transformations.
|
||
"""
|
||
if resample not in _ROTATE_SAFE_FILTERS:
|
||
return Image.Resampling.BICUBIC
|
||
return resample
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Tensor <-> PIL conversion helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def tensor_to_pil(tensor: torch.Tensor) -> Image.Image:
|
||
"""
|
||
Convert a single ComfyUI IMAGE tensor (H, W, C) in [0, 1] float32
|
||
to a PIL RGB Image.
|
||
"""
|
||
# Clamp to valid range and convert to uint8
|
||
arr = (tensor.clamp(0.0, 1.0).cpu().numpy() * 255.0).astype(np.uint8)
|
||
return Image.fromarray(arr, mode="RGB")
|
||
|
||
|
||
def pil_to_tensor(image: Image.Image) -> torch.Tensor:
|
||
"""
|
||
Convert a PIL RGB Image to a ComfyUI IMAGE tensor (H, W, C) in [0, 1] float32.
|
||
"""
|
||
arr = np.array(image.convert("RGB")).astype(np.float32) / 255.0
|
||
return torch.from_numpy(arr)
|
||
|
||
|
||
def mask_tensor_to_pil(tensor: torch.Tensor) -> Image.Image:
|
||
"""
|
||
Convert a single ComfyUI MASK tensor (H, W) in [0, 1] float32
|
||
to a PIL grayscale Image.
|
||
"""
|
||
arr = (tensor.clamp(0.0, 1.0).cpu().numpy() * 255.0).astype(np.uint8)
|
||
return Image.fromarray(arr, mode="L")
|
||
|
||
|
||
def pil_to_mask_tensor(image: Image.Image) -> torch.Tensor:
|
||
"""
|
||
Convert a PIL grayscale Image to a ComfyUI MASK tensor (H, W) in [0, 1] float32.
|
||
"""
|
||
arr = np.array(image.convert("L")).astype(np.float32) / 255.0
|
||
return torch.from_numpy(arr)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Rotation helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _get_resample_filter(method_name: str) -> Image.Resampling:
|
||
"""Return the PIL resampling filter for the given method name."""
|
||
return INTERPOLATION_METHODS.get(method_name, INTERPOLATION_METHODS[DEFAULT_INTERPOLATION])
|
||
|
||
|
||
def rotate_pil_image(
|
||
image: Image.Image,
|
||
angle: float,
|
||
resample: Image.Resampling,
|
||
expand: bool = True,
|
||
fill_color: Tuple[int, int, int] = (0, 0, 0),
|
||
) -> Image.Image:
|
||
"""
|
||
Rotate a PIL RGB image by *angle* degrees (counter-clockwise, PIL convention).
|
||
|
||
Parameters
|
||
----------
|
||
image:
|
||
Source PIL Image (RGB).
|
||
angle:
|
||
Rotation angle in degrees. Positive = counter-clockwise.
|
||
resample:
|
||
PIL resampling filter (e.g. Image.Resampling.BICUBIC).
|
||
Note: PIL's rotate() only supports NEAREST, BILINEAR, and BICUBIC.
|
||
If LANCZOS is passed it is automatically promoted to BICUBIC.
|
||
expand:
|
||
If True the output canvas is expanded to fit the full rotated image.
|
||
If False the canvas stays the same size (content may be clipped).
|
||
fill_color:
|
||
RGB tuple used to fill the new areas introduced by rotation.
|
||
|
||
Returns
|
||
-------
|
||
PIL Image (RGB) after rotation.
|
||
"""
|
||
return image.rotate(
|
||
angle,
|
||
resample=_rotate_resample(resample),
|
||
expand=expand,
|
||
fillcolor=fill_color,
|
||
)
|
||
|
||
|
||
def rotate_pil_mask(
|
||
mask: Image.Image,
|
||
angle: float,
|
||
resample: Image.Resampling,
|
||
expand: bool = True,
|
||
) -> Image.Image:
|
||
"""
|
||
Rotate a PIL grayscale mask by *angle* degrees.
|
||
|
||
New areas introduced by rotation are filled with black (0).
|
||
Masks always use NEAREST or the provided resample filter; for masks
|
||
NEAREST is often preferred to keep binary values crisp, but we respect
|
||
the caller's choice so that the mask matches the image geometry exactly.
|
||
|
||
Parameters
|
||
----------
|
||
mask:
|
||
Source PIL Image (L / grayscale).
|
||
angle:
|
||
Rotation angle in degrees. Positive = counter-clockwise.
|
||
resample:
|
||
PIL resampling filter.
|
||
Note: PIL's rotate() only supports NEAREST, BILINEAR, and BICUBIC.
|
||
If LANCZOS is passed it is automatically promoted to BICUBIC.
|
||
expand:
|
||
Same semantics as for rotate_pil_image.
|
||
|
||
Returns
|
||
-------
|
||
PIL Image (L) after rotation.
|
||
"""
|
||
return mask.rotate(
|
||
angle,
|
||
resample=_rotate_resample(resample),
|
||
expand=expand,
|
||
fillcolor=0,
|
||
)
|
||
|
||
|
||
def center_crop(image: Image.Image, target_w: int, target_h: int) -> Image.Image:
|
||
"""
|
||
Crop *image* to *target_w* × *target_h* from the centre.
|
||
|
||
If the image is smaller than the target in any dimension the image is
|
||
returned unchanged (no padding is applied — callers should ensure the
|
||
rotated canvas is always at least as large as the original).
|
||
"""
|
||
w, h = image.size
|
||
left = max(0, (w - target_w) // 2)
|
||
top = max(0, (h - target_h) // 2)
|
||
right = left + min(target_w, w)
|
||
bottom = top + min(target_h, h)
|
||
return image.crop((left, top, right, bottom))
|
||
|
||
|
||
def compute_rotated_canvas_size(
|
||
original_w: int, original_h: int, angle_deg: float
|
||
) -> Tuple[int, int]:
|
||
"""
|
||
Compute the canvas size required to contain the full rotated image.
|
||
|
||
Returns (new_width, new_height) as integers.
|
||
"""
|
||
angle_rad = math.radians(angle_deg)
|
||
cos_a = abs(math.cos(angle_rad))
|
||
sin_a = abs(math.sin(angle_rad))
|
||
new_w = int(math.ceil(original_w * cos_a + original_h * sin_a))
|
||
new_h = int(math.ceil(original_w * sin_a + original_h * cos_a))
|
||
return new_w, new_h
|