MachinePainting
21fbdb901e
offset_y was following the raw pixel convention where Y grows downward, so dragging the slider right pushed the image down. Negated it in build_matrix so both placement sliders now read the same way: drag right for right/up, drag left for left/down. One edit covers both nodes, since Image Transform Pro and Remove Background Pro share build_matrix. This is the opposite sign to Photoshop's transform Y field, which is a deliberate choice: on a centre-zero slider "more" reading as "up" is more intuitive than matching pixel coordinates. Category is now "MachinePaintingNodes/Rotate & Scale" - a single subcategory using "&" rather than "/", which ComfyUI would have split into nested submenus. Keeps the node one level deep like the rest of the pack. README heading updated to match.
190 lines
6.8 KiB
Python
190 lines
6.8 KiB
Python
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
import folder_paths
|
|
import os
|
|
|
|
from .transform_utils import (
|
|
build_matrix,
|
|
coverage_mask,
|
|
fit_scale,
|
|
normalize_mask,
|
|
snap_angle,
|
|
transparency_grid,
|
|
warp_image,
|
|
warp_single,
|
|
zoom_to_scale,
|
|
)
|
|
|
|
|
|
class ImageTransformPro:
|
|
"""
|
|
Rotate, scale and reposition an image inside a definable canvas.
|
|
|
|
Rotation reads counter-clockwise to the left of centre and clockwise to the
|
|
right. The "fixed" method snaps the angle to 45 degree increments. Zoom is
|
|
centred on 0 (1.0x) and exponential: -100 = 0.25x, -50 = 0.5x, +50 = 2x,
|
|
+100 = 4x. An optional mask input rides the exact same matrix so it stays
|
|
registered to the image, and cuts the image out: whatever the mask leaves
|
|
out is filled with the selected background. With no mask supplied, the
|
|
transformed image extent becomes the mask.
|
|
"""
|
|
|
|
# ComfyUI reads DESCRIPTION (not the docstring) for the node tooltip and
|
|
# indexes it for node search, so name the verbs the node is looked up by.
|
|
DESCRIPTION = (
|
|
"Rotate, scale, zoom and move an image inside a definable canvas. "
|
|
"Free or fixed 45 degree rotation, centre-zero zoom and XY placement. "
|
|
"An optional mask is transformed with the image and cuts it out over a "
|
|
"transparent, black or white background."
|
|
)
|
|
|
|
ROTATION_METHODS = ["free", "fixed"]
|
|
BACKGROUNDS = ["transparent", "black", "white"]
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
},
|
|
"optional": {
|
|
"mask": ("MASK",),
|
|
|
|
# --- Canvas (0 = match the incoming image) ---
|
|
"canvas_width": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 8192,
|
|
"step": 8
|
|
}),
|
|
"canvas_height": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 8192,
|
|
"step": 8
|
|
}),
|
|
"background": (cls.BACKGROUNDS, {"default": "transparent"}),
|
|
|
|
# --- Rotation ---
|
|
"rotation_method": (cls.ROTATION_METHODS, {"default": "free"}),
|
|
"rotation_angle": ("FLOAT", {
|
|
"default": 0.0,
|
|
"min": -180.0,
|
|
"max": 180.0,
|
|
"step": 0.1,
|
|
"display": "slider"
|
|
}),
|
|
|
|
# --- Zoom ---
|
|
"zoom": ("FLOAT", {
|
|
"default": 0.0,
|
|
"min": -100.0,
|
|
"max": 100.0,
|
|
"step": 0.1,
|
|
"display": "slider"
|
|
}),
|
|
|
|
# --- Placement (percent of canvas) ---
|
|
"offset_x": ("FLOAT", {
|
|
"default": 0.0,
|
|
"min": -100.0,
|
|
"max": 100.0,
|
|
"step": 0.1,
|
|
"display": "slider"
|
|
}),
|
|
"offset_y": ("FLOAT", {
|
|
"default": 0.0,
|
|
"min": -100.0,
|
|
"max": 100.0,
|
|
"step": 0.1,
|
|
"display": "slider"
|
|
}),
|
|
},
|
|
"hidden": {
|
|
"unique_id": "UNIQUE_ID",
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK", "IMAGE")
|
|
RETURN_NAMES = ("image", "mask", "mask_bw")
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "rotate_image"
|
|
CATEGORY = "MachinePaintingNodes/Rotate & Scale"
|
|
|
|
def rotate_image(self, image, mask=None, canvas_width=0, canvas_height=0,
|
|
background="transparent", rotation_method="free",
|
|
rotation_angle=0.0, zoom=0.0, offset_x=0.0, offset_y=0.0,
|
|
unique_id=None):
|
|
|
|
img = image[0].cpu().numpy().astype(np.float32)
|
|
src_h, src_w = img.shape[:2]
|
|
|
|
# Canvas: 0 means "match the incoming image"
|
|
dst_w = int(canvas_width) if canvas_width > 0 else src_w
|
|
dst_h = int(canvas_height) if canvas_height > 0 else src_h
|
|
dst_shape = (dst_h, dst_w)
|
|
|
|
angle = snap_angle(rotation_angle, rotation_method)
|
|
|
|
# Fit the source inside the canvas first, then apply zoom on top, so
|
|
# zoom 0 always shows the whole image whatever the canvas size is.
|
|
scale = fit_scale(img.shape, dst_shape) * zoom_to_scale(zoom)
|
|
|
|
M = build_matrix(img.shape, dst_shape, angle, scale, offset_x, offset_y)
|
|
|
|
warped = warp_image(img, M, dst_shape)
|
|
coverage = coverage_mask(img.shape, M, dst_shape)
|
|
|
|
# An incoming mask rides the same matrix so it stays aligned, and it
|
|
# decides what stays visible in the image. Without one, the transformed
|
|
# image extent becomes the mask.
|
|
if mask is not None:
|
|
mask_np = normalize_mask(mask, (src_h, src_w))
|
|
# Multiply by coverage so the mask can never claim area outside the
|
|
# transformed image, even when it is white right to its own edges.
|
|
out_mask = warp_single(mask_np, M, dst_shape) * coverage
|
|
else:
|
|
out_mask = coverage
|
|
|
|
# Fill everything that isn't visible. "transparent" is carried by the
|
|
# mask output, since a ComfyUI IMAGE tensor has no alpha channel.
|
|
if background == "white":
|
|
bg = np.ones((dst_h, dst_w, 3), dtype=np.float32)
|
|
else:
|
|
bg = np.zeros((dst_h, dst_w, 3), dtype=np.float32)
|
|
|
|
vis_3ch = np.expand_dims(out_mask, axis=-1)
|
|
result = warped * vis_3ch + bg * (1.0 - vis_3ch)
|
|
result = np.clip(result, 0.0, 1.0)
|
|
|
|
# Build outputs
|
|
result_tensor = torch.from_numpy(result).unsqueeze(0)
|
|
mask_tensor = torch.from_numpy(out_mask).unsqueeze(0)
|
|
|
|
mask_bw = np.stack([out_mask, out_mask, out_mask], axis=2)
|
|
mask_bw_tensor = torch.from_numpy(mask_bw).unsqueeze(0)
|
|
|
|
# Preview
|
|
if background == "transparent":
|
|
grid = transparency_grid(result.shape)
|
|
preview = (result * 255.0) * vis_3ch + grid * (1.0 - vis_3ch)
|
|
preview_img = np.clip(preview, 0, 255).astype(np.uint8)
|
|
else:
|
|
preview_img = (result * 255.0).astype(np.uint8)
|
|
|
|
preview_results = self._save_preview(preview_img, unique_id)
|
|
|
|
return {
|
|
"ui": {"images": preview_results},
|
|
"result": (result_tensor, mask_tensor, mask_bw_tensor)
|
|
}
|
|
|
|
def _save_preview(self, img_np, unique_id):
|
|
temp_dir = folder_paths.get_temp_directory()
|
|
filename = f"image_transform_{unique_id}.png"
|
|
filepath = os.path.join(temp_dir, filename)
|
|
img_pil = Image.fromarray(img_np, mode='RGB')
|
|
img_pil.save(filepath)
|
|
return [{"filename": filename, "subfolder": "", "type": "temp"}]
|