268 lines
11 KiB
Python
268 lines
11 KiB
Python
import torch
|
||
import random
|
||
import math
|
||
import torch.nn.functional as F
|
||
|
||
class RandomAspectRatioMask:
|
||
def __init__(self):
|
||
self._current_aspect_ratio = "16:9"
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
ar_presets = ["random", "16:9", "9:16", "3:2", "2:3", "4:5", "5:4", "1:1"]
|
||
placements = [
|
||
"center", "random",
|
||
"left", "right", "up", "down",
|
||
"top-left", "top-mid", "top-right",
|
||
"mid-left", "mid-right",
|
||
"bottom-left", "bottom-mid", "bottom-right"
|
||
]
|
||
scales = ["random", "50", "60", "70", "80", "90", "100"]
|
||
rotations = ["random", "0", "30", "60", "90", "120", "150", "180", "210", "240", "270", "300", "330"]
|
||
bg_colors = ["white", "black", "grey"]
|
||
return {
|
||
"required": {
|
||
"image": ("IMAGE",),
|
||
"aspect_ratio": (ar_presets, {"default": "16:9"}),
|
||
"placement": (placements, {"default": "center"}),
|
||
"scale_pct": (scales, {"default": "100"}),
|
||
"rotation": (rotations, {"default": "0"}),
|
||
"bg_color": (bg_colors, {"default": "white"}),
|
||
"seed": ("INT", {"default": 0, "min": 0, "max": 2**31-1}),
|
||
"feathering": ("INT", {"default": 40, "min": 0, "max": 1024}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE", "MASK", "STRING")
|
||
RETURN_NAMES = ("Image", "Mask", "Params")
|
||
FUNCTION = "process"
|
||
CATEGORY = "va1"
|
||
|
||
def process(self, image, aspect_ratio, placement, scale_pct, rotation, bg_color, seed, feathering):
|
||
# ── 1) Handle aspect_ratio (random if chosen) ─────────────────────────────
|
||
presets = ["16:9", "9:16", "3:2", "2:3", "4:5", "5:4", "1:1"]
|
||
if aspect_ratio == "random":
|
||
aspect_ratio = random.choice(presets)
|
||
self._current_aspect_ratio = aspect_ratio
|
||
|
||
# ── 1.5) Handle rotation (random if chosen) ──────────────────────────────
|
||
rotation_options = [150, 180, 210, 240, 270, 300, 330, 0, 30, 60, 90, 120]
|
||
if rotation == "random":
|
||
rotation_deg = random.choice(rotation_options)
|
||
else:
|
||
rotation_deg = int(rotation)
|
||
|
||
# Determine background color value early for rotation
|
||
bg_value = 1.0 if bg_color == "white" else (0.0 if bg_color == "black" else 0.5)
|
||
|
||
# Apply rotation to input image if needed
|
||
if rotation_deg != 0:
|
||
image = self._rotate_image(image, rotation_deg, bg_value)
|
||
|
||
# ── 2) Compute canvas size from ORIGINAL image dims ────────────────────────
|
||
b0, h0, w0, c0 = image.size()
|
||
try:
|
||
wf, hf = map(float, aspect_ratio.split(":"))
|
||
target_ar = wf / hf
|
||
except:
|
||
target_ar = 1.0
|
||
orig_ar = w0 / h0
|
||
|
||
if abs(orig_ar - target_ar) < 1e-6:
|
||
new_h, new_w = h0, w0
|
||
elif target_ar > orig_ar:
|
||
new_w, new_h = int(target_ar * h0), h0
|
||
else:
|
||
new_h, new_w = int(w0 / target_ar), w0
|
||
|
||
# Round to multiples of 8
|
||
new_w = ((new_w + 7) // 8) * 8
|
||
new_h = ((new_h + 7) // 8) * 8
|
||
|
||
# ── 3) Pick scale (after canvas dims)—100% disallowed for 1:1 ──────────────
|
||
all_scales = ["50", "60", "70", "80", "90", "100"]
|
||
valid_scales = all_scales[:-2] if aspect_ratio == "1:1" else all_scales
|
||
|
||
if scale_pct == "random":
|
||
chosen_scale = random.choice(valid_scales)
|
||
else:
|
||
chosen_scale = scale_pct if scale_pct in valid_scales else valid_scales[-1]
|
||
|
||
scale = int(chosen_scale) / 100.0
|
||
|
||
# Store selected parameters for output
|
||
selected_scale = chosen_scale
|
||
selected_rotation = str(rotation_deg)
|
||
|
||
# ── 4) Resize the **original image** now, so canvas stays fixed ───────────
|
||
if scale != 1.0:
|
||
nh0 = max(1, int(h0 * scale))
|
||
nw0 = max(1, int(w0 * scale))
|
||
# B×H×W×C → B×C×H×W for F.interpolate
|
||
img_bc = image.permute(0, 3, 1, 2)
|
||
resized = F.interpolate(
|
||
img_bc,
|
||
size=(nh0, nw0),
|
||
mode='bilinear',
|
||
align_corners=False
|
||
)
|
||
image = resized.permute(0, 2, 3, 1)
|
||
b, h, w, c = b0, nh0, nw0, c0
|
||
else:
|
||
b, h, w, c = b0, h0, w0, c0
|
||
|
||
# ── 5) Determine placement offsets ─────────────────────────────────────────
|
||
# Orientation check for directional fallbacks
|
||
horizontal = (target_ar >= 1.0)
|
||
vertical = (target_ar < 1.0)
|
||
effective = placement
|
||
if placement in ("left", "right") and not horizontal:
|
||
effective = "center"
|
||
if placement in ("up", "down") and not vertical:
|
||
effective = "center"
|
||
|
||
mid_x = (new_w - w) // 2
|
||
mid_y = (new_h - h) // 2
|
||
|
||
pos_map = {
|
||
"center": (mid_x, mid_y),
|
||
"random": (random.randint(0, new_w - w),
|
||
random.randint(0, new_h - h)),
|
||
"left": (0, mid_y),
|
||
"right": (new_w - w, mid_y),
|
||
"up": (mid_x, 0),
|
||
"down": (mid_x, new_h - h),
|
||
"top-left": (0, 0),
|
||
"top-mid": (mid_x, 0),
|
||
"top-right": (new_w - w, 0),
|
||
"mid-left": (0, mid_y),
|
||
"mid-right": (new_w - w, mid_y),
|
||
"bottom-left": (0, new_h - h),
|
||
"bottom-mid": (mid_x, new_h - h),
|
||
"bottom-right":(new_w - w, new_h - h),
|
||
}
|
||
|
||
pad_l, pad_t = pos_map.get(effective, pos_map["center"])
|
||
pad_r = new_w - w - pad_l
|
||
pad_b = new_h - h - pad_t
|
||
|
||
# Store selected placement for output
|
||
selected_placement = effective
|
||
|
||
# ── 6) Build padded image ─────────────────────────────────────────────────
|
||
new_image = torch.full(
|
||
(b, new_h, new_w, c),
|
||
bg_value,
|
||
dtype=image.dtype,
|
||
device=image.device
|
||
)
|
||
new_image[:, pad_t:pad_t+h, pad_l:pad_l+w, :] = image
|
||
|
||
# ── 7) Feathered mask ─────────────────────────────────────────────────────
|
||
mask = torch.ones((new_h, new_w),
|
||
dtype=image.dtype,
|
||
device=image.device)
|
||
inner = torch.zeros((h, w),
|
||
dtype=image.dtype,
|
||
device=image.device)
|
||
if feathering > 0 and feathering * 2 < h and feathering * 2 < w:
|
||
for i in range(h):
|
||
for j in range(w):
|
||
dt = i if pad_t > 0 else h
|
||
db = (h - 1 - i) if pad_b > 0 else h
|
||
dl = j if pad_l > 0 else w
|
||
dr = (w - 1 - j) if pad_r > 0 else w
|
||
d = min(dt, db, dl, dr)
|
||
if d < feathering:
|
||
v = (feathering - d) / feathering
|
||
inner[i, j] = v * v
|
||
|
||
mask[pad_t:pad_t+h, pad_l:pad_l+w] = inner
|
||
|
||
# Create parameters string
|
||
params_str = f"AR: {aspect_ratio}, Placement: {selected_placement}, Scale: {selected_scale}%, Rotation: {selected_rotation}°"
|
||
|
||
return (new_image, mask.unsqueeze(0), params_str)
|
||
|
||
def _rotate_image(self, image, degrees, bg_value=0.5):
|
||
"""Rotate image by given degrees using bilinear interpolation with specified background color."""
|
||
if degrees == 0:
|
||
return image
|
||
|
||
b, h, w, c = image.shape
|
||
|
||
# Convert degrees to radians
|
||
angle_rad = math.radians(degrees)
|
||
cos_a = math.cos(angle_rad)
|
||
sin_a = math.sin(angle_rad)
|
||
|
||
# Calculate new dimensions after rotation
|
||
new_h = int(abs(h * cos_a) + abs(w * sin_a))
|
||
new_w = int(abs(w * cos_a) + abs(h * sin_a))
|
||
|
||
# Create coordinate grids for the output image
|
||
device = image.device
|
||
dtype = image.dtype
|
||
|
||
# Create background canvas with specified color
|
||
rotated_image = torch.full(
|
||
(b, new_h, new_w, c),
|
||
bg_value,
|
||
dtype=dtype,
|
||
device=device
|
||
)
|
||
|
||
# Center coordinates
|
||
center_h = (h - 1) / 2
|
||
center_w = (w - 1) / 2
|
||
new_center_h = (new_h - 1) / 2
|
||
new_center_w = (new_w - 1) / 2
|
||
|
||
# Create meshgrid for new image coordinates
|
||
y_coords = torch.arange(new_h, dtype=dtype, device=device)
|
||
x_coords = torch.arange(new_w, dtype=dtype, device=device)
|
||
yy, xx = torch.meshgrid(y_coords, x_coords, indexing='ij')
|
||
|
||
# Translate to center, rotate, then translate back
|
||
yy_centered = yy - new_center_h
|
||
xx_centered = xx - new_center_w
|
||
|
||
# Apply inverse rotation to find source coordinates
|
||
yy_rot = yy_centered * cos_a + xx_centered * sin_a
|
||
xx_rot = -yy_centered * sin_a + xx_centered * cos_a
|
||
|
||
# Translate back to original coordinate system
|
||
yy_orig = yy_rot + center_h
|
||
xx_orig = xx_rot + center_w
|
||
|
||
# Normalize coordinates to [-1, 1] for grid_sample
|
||
yy_norm = 2.0 * yy_orig / (h - 1) - 1.0
|
||
xx_norm = 2.0 * xx_orig / (w - 1) - 1.0
|
||
|
||
# Stack coordinates for grid_sample (N, H, W, 2)
|
||
grid = torch.stack([xx_norm, yy_norm], dim=-1).unsqueeze(0)
|
||
|
||
# Reshape image for grid_sample: (B, C, H, W)
|
||
image_permuted = image.permute(0, 3, 1, 2)
|
||
|
||
# Apply rotation using grid_sample with reflection padding
|
||
rotated_temp = F.grid_sample(
|
||
image_permuted,
|
||
grid.repeat(b, 1, 1, 1),
|
||
mode='bilinear',
|
||
padding_mode='reflection',
|
||
align_corners=True
|
||
)
|
||
|
||
# Reshape back to (B, H, W, C)
|
||
rotated_temp = rotated_temp.permute(0, 2, 3, 1)
|
||
|
||
# Create a mask for valid pixels (within original image bounds)
|
||
mask = ((yy_orig >= 0) & (yy_orig < h) & (xx_orig >= 0) & (xx_orig < w)).float()
|
||
mask = mask.unsqueeze(0).unsqueeze(-1) # Add batch and channel dimensions
|
||
|
||
# Blend the rotated image with the background using the mask
|
||
rotated_image = rotated_image * (1 - mask) + rotated_temp * mask
|
||
|
||
return rotated_image
|