145 lines
4.4 KiB
Python
145 lines
4.4 KiB
Python
"""
|
|
RotateImage — ComfyUI custom node.
|
|
|
|
Rotates an IMAGE (and optional MASK) by a user-specified angle, expanding
|
|
the canvas and filling new areas with black. Outputs the rotated image,
|
|
the rotated mask, and a ROTATION_INFO dict that the paired
|
|
RestoreRotatedImage node needs to reverse the transformation.
|
|
"""
|
|
|
|
from typing import Optional, Tuple
|
|
|
|
import torch
|
|
|
|
from ri_utils.rotation import (
|
|
DEFAULT_INTERPOLATION,
|
|
INTERPOLATION_METHODS,
|
|
center_crop,
|
|
mask_tensor_to_pil,
|
|
pil_to_mask_tensor,
|
|
pil_to_tensor,
|
|
rotate_pil_image,
|
|
rotate_pil_mask,
|
|
tensor_to_pil,
|
|
)
|
|
|
|
|
|
class RotateImage:
|
|
"""
|
|
Rotate an image (and optional mask) by a given angle.
|
|
|
|
The canvas is expanded so that no content is clipped; new areas are
|
|
filled with black. A ROTATION_INFO dict is emitted so that the
|
|
RestoreRotatedImage node can exactly reverse the transformation.
|
|
"""
|
|
|
|
CATEGORY = "image/transform"
|
|
FUNCTION = "rotate"
|
|
RETURN_TYPES = ("IMAGE", "MASK", "ROTATION_INFO")
|
|
RETURN_NAMES = ("image", "mask", "rotation_info")
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls): # noqa: N802 — ComfyUI naming convention
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"angle": (
|
|
"FLOAT",
|
|
{
|
|
"default": 0.0,
|
|
"min": -360.0,
|
|
"max": 360.0,
|
|
"step": 0.1,
|
|
"display": "number",
|
|
},
|
|
),
|
|
"interpolation": (
|
|
list(INTERPOLATION_METHODS.keys()),
|
|
{"default": DEFAULT_INTERPOLATION},
|
|
), # COMBO: first element is the list of choices
|
|
},
|
|
"optional": {
|
|
"mask": ("MASK",),
|
|
},
|
|
}
|
|
|
|
# ------------------------------------------------------------------
|
|
# Main entry point
|
|
# ------------------------------------------------------------------
|
|
|
|
def rotate(
|
|
self,
|
|
image: torch.Tensor,
|
|
angle: float,
|
|
interpolation: str,
|
|
mask: Optional[torch.Tensor] = None,
|
|
) -> Tuple[torch.Tensor, torch.Tensor, dict]:
|
|
"""
|
|
Parameters
|
|
----------
|
|
image:
|
|
ComfyUI IMAGE tensor, shape (B, H, W, C), float32 in [0, 1].
|
|
angle:
|
|
Rotation angle in degrees. Positive = counter-clockwise
|
|
(PIL convention).
|
|
interpolation:
|
|
One of the keys in INTERPOLATION_METHODS.
|
|
mask:
|
|
Optional ComfyUI MASK tensor, shape (B, H, W), float32 in [0, 1].
|
|
|
|
Returns
|
|
-------
|
|
(rotated_image, rotated_mask, rotation_info)
|
|
"""
|
|
resample = INTERPOLATION_METHODS.get(
|
|
interpolation, INTERPOLATION_METHODS[DEFAULT_INTERPOLATION]
|
|
)
|
|
|
|
# Original spatial dimensions (from the first image in the batch)
|
|
_, orig_h, orig_w, _ = image.shape
|
|
|
|
rotated_images = []
|
|
for i in range(image.shape[0]):
|
|
pil_img = tensor_to_pil(image[i])
|
|
rotated = rotate_pil_image(
|
|
pil_img,
|
|
angle=angle,
|
|
resample=resample,
|
|
expand=True,
|
|
fill_color=(0, 0, 0),
|
|
)
|
|
rotated_images.append(pil_to_tensor(rotated))
|
|
|
|
# Stack back to (B, H', W', C)
|
|
out_image = torch.stack(rotated_images, dim=0)
|
|
|
|
# --- Mask ---
|
|
if mask is not None:
|
|
rotated_masks = []
|
|
# Mask may have a different batch size; handle gracefully
|
|
for i in range(mask.shape[0]):
|
|
pil_mask = mask_tensor_to_pil(mask[i])
|
|
rotated_m = rotate_pil_mask(
|
|
pil_mask,
|
|
angle=angle,
|
|
resample=resample,
|
|
expand=True,
|
|
)
|
|
rotated_masks.append(pil_to_mask_tensor(rotated_m))
|
|
out_mask = torch.stack(rotated_masks, dim=0)
|
|
else:
|
|
# Emit a black (all-zero) mask matching the rotated image size
|
|
_, rh, rw, _ = out_image.shape
|
|
out_mask = torch.zeros(
|
|
(image.shape[0], rh, rw), dtype=torch.float32
|
|
)
|
|
|
|
rotation_info = {
|
|
"angle": angle,
|
|
"original_width": orig_w,
|
|
"original_height": orig_h,
|
|
"interpolation": interpolation,
|
|
}
|
|
|
|
return (out_image, out_mask, rotation_info)
|