RyanOnTheInside
@@ -1 +1,2 @@
|
||||
# This file intentionally left blank
|
||||
|
||||
@@ -0,0 +1,420 @@
|
||||
"""
|
||||
Optical flow nodes for generating and applying flow-based warping.
|
||||
|
||||
These nodes work with optical flow to create more advanced warping effects.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import math
|
||||
|
||||
from ....src.realtimenodes.control_base import ControlNodeBase
|
||||
from ....src.utils.general import AlwaysEqualProxy
|
||||
from ....src.utils.image import gaussian_blur_2d, flow_to_rgb
|
||||
from ....src.utils.realtime_flownets import RealTimeFlowNet
|
||||
|
||||
|
||||
class OpticalFlowNode(ControlNodeBase):
|
||||
"""
|
||||
Generate optical flow between images.
|
||||
|
||||
Computes the motion flow field between two images.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
inputs = super().INPUT_TYPES()
|
||||
inputs["required"].update({
|
||||
"previous_image": ("IMAGE", {"tooltip": "Previous frame image"}),
|
||||
"current_image": ("IMAGE", {"tooltip": "Current frame image"}),
|
||||
"flow_scale": (
|
||||
"FLOAT",
|
||||
{"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1,
|
||||
"tooltip": "Scale factor to apply to flow"}
|
||||
),
|
||||
"smoothness": (
|
||||
"FLOAT",
|
||||
{"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1,
|
||||
"tooltip": "Gaussian blur sigma to smooth flow field"}
|
||||
),
|
||||
"visualize": (
|
||||
"BOOLEAN",
|
||||
{"default": False, "tooltip": "Output flow visualization as RGB image"}
|
||||
),
|
||||
})
|
||||
return inputs
|
||||
|
||||
RETURN_TYPES = ("FLOW_FIELD", "IMAGE")
|
||||
RETURN_NAMES = ("flow", "flow_visualization")
|
||||
FUNCTION = "update"
|
||||
CATEGORY = "Realtime Nodes/control/flow"
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.flow_net = None
|
||||
|
||||
def update(
|
||||
self,
|
||||
previous_image,
|
||||
current_image,
|
||||
flow_scale=1.0,
|
||||
smoothness=1.0,
|
||||
visualize=False,
|
||||
always_execute=True,
|
||||
unique_id=None
|
||||
):
|
||||
"""
|
||||
Generate optical flow between two images.
|
||||
|
||||
Args:
|
||||
previous_image: Previous frame tensor [B, H, W, C]
|
||||
current_image: Current frame tensor [B, H, W, C]
|
||||
flow_scale: Factor to scale the flow magnitude
|
||||
smoothness: Gaussian blur sigma to apply to the flow field
|
||||
visualize: Whether to output a visualization of the flow field
|
||||
always_execute: Whether to always execute (from ControlNodeBase)
|
||||
unique_id: Unique ID for state management (from ControlNodeBase)
|
||||
|
||||
Returns:
|
||||
Flow field tensor [B, 2, H, W] and optional visualization
|
||||
"""
|
||||
# Initialize flow network on first use
|
||||
if self.flow_net is None:
|
||||
self.flow_net = RealTimeFlowNet()
|
||||
self.flow_net.eval()
|
||||
|
||||
# Ensure inputs are on the same device
|
||||
device = previous_image.device
|
||||
self.flow_net = self.flow_net.to(device)
|
||||
|
||||
# Convert to grayscale and ensure proper format for flow estimation
|
||||
# Input is BHWC, convert to BCHW for CNN
|
||||
prev_gray = torch.mean(previous_image, dim=-1, keepdim=True).permute(0, 3, 1, 2)
|
||||
curr_gray = torch.mean(current_image, dim=-1, keepdim=True).permute(0, 3, 1, 2)
|
||||
|
||||
# Stack frames for flow estimation
|
||||
stacked = torch.cat([prev_gray, curr_gray], dim=1)
|
||||
|
||||
# Estimate flow
|
||||
with torch.no_grad():
|
||||
flow = self.flow_net(stacked)
|
||||
|
||||
# Scale the flow
|
||||
flow = flow * flow_scale
|
||||
|
||||
# Apply Gaussian blur if smoothness > 0
|
||||
if smoothness > 0:
|
||||
flow = gaussian_blur_2d(flow, kernel_size=9, sigma=smoothness)
|
||||
|
||||
# Create flow visualization if requested
|
||||
if visualize:
|
||||
flow_rgb = flow_to_rgb(flow)
|
||||
# Convert back to BHWC for ComfyUI
|
||||
flow_vis = flow_rgb.permute(0, 2, 3, 1)
|
||||
else:
|
||||
# Create empty tensor for visualization
|
||||
flow_vis = torch.zeros_like(previous_image)
|
||||
|
||||
# Print flow statistics for debugging
|
||||
with torch.no_grad():
|
||||
flow_abs_mean = torch.abs(flow).mean().item()
|
||||
flow_min = flow.min().item()
|
||||
flow_max = flow.max().item()
|
||||
print(f"Flow stats - mean abs: {flow_abs_mean:.4f}, min: {flow_min:.4f}, max: {flow_max:.4f}")
|
||||
|
||||
return (flow, flow_vis)
|
||||
|
||||
|
||||
class FlowWarpNode(ControlNodeBase):
|
||||
"""
|
||||
Apply optical flow field to warp images or latents.
|
||||
|
||||
Uses flow fields for advanced warping effects.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
inputs = super().INPUT_TYPES()
|
||||
inputs["required"].update({
|
||||
"data_type": (
|
||||
["IMAGE", "LATENT"],
|
||||
{"default": "IMAGE", "tooltip": "Type of data to warp"}
|
||||
),
|
||||
"data": (
|
||||
AlwaysEqualProxy("*"),
|
||||
{"tooltip": "Image or latent to warp with flow field"}
|
||||
),
|
||||
"flow_field": (
|
||||
"FLOW_FIELD",
|
||||
{"tooltip": "Flow field tensor [B, 2, H, W]"}
|
||||
),
|
||||
"strength": (
|
||||
"FLOAT",
|
||||
{"default": 5.0, "min": 0.0, "max": 50.0, "step": 0.1,
|
||||
"tooltip": "Strength of the flow warping effect (try higher values if no effect is visible)"}
|
||||
),
|
||||
"mode": (
|
||||
["bilinear", "nearest", "bicubic"],
|
||||
{"default": "bilinear", "tooltip": "Interpolation mode for sampling"}
|
||||
),
|
||||
"padding_mode": (
|
||||
["zeros", "border", "reflection"],
|
||||
{"default": "border", "tooltip": "How to handle pixels outside boundaries"}
|
||||
),
|
||||
"align_corners": (
|
||||
"BOOLEAN",
|
||||
{"default": True, "tooltip": "Whether to align corners in sampling grid"}
|
||||
),
|
||||
"debug_mode": (
|
||||
["none", "flow", "grid", "test_pattern"],
|
||||
{"default": "none", "tooltip": "Visualization mode for debugging"}
|
||||
),
|
||||
})
|
||||
return inputs
|
||||
|
||||
RETURN_TYPES = (AlwaysEqualProxy("*"), "IMAGE")
|
||||
RETURN_NAMES = ("warped_data", "debug_visualization")
|
||||
FUNCTION = "update"
|
||||
CATEGORY = "Realtime Nodes/control/flow"
|
||||
|
||||
def create_test_pattern(self, height, width, device):
|
||||
"""Create a test pattern grid image to visualize warping effects"""
|
||||
# Create checkerboard pattern
|
||||
y, x = torch.meshgrid(torch.linspace(0, height-1, height, device=device),
|
||||
torch.linspace(0, width-1, width, device=device),
|
||||
indexing="ij")
|
||||
|
||||
# Create grid lines
|
||||
grid_size = 32
|
||||
grid_x = (x % grid_size < 2).float()
|
||||
grid_y = (y % grid_size < 2).float()
|
||||
grid = grid_x + grid_y
|
||||
|
||||
# Create circles
|
||||
centers = [(height//4, width//4), (height//4, 3*width//4),
|
||||
(3*height//4, width//4), (3*height//4, 3*width//4)]
|
||||
circles = torch.zeros_like(grid)
|
||||
for cy, cx in centers:
|
||||
dist = torch.sqrt((y - cy)**2 + (x - cx)**2)
|
||||
circles += (dist < 30).float()
|
||||
|
||||
# Combine patterns
|
||||
r = grid
|
||||
g = (x / width)
|
||||
b = (y / height)
|
||||
|
||||
# Add circles
|
||||
r = torch.clamp(r + circles, 0, 1)
|
||||
|
||||
# Stack channels
|
||||
pattern = torch.stack([r, g, b], dim=-1)
|
||||
|
||||
# Add batch dimension
|
||||
pattern = pattern.unsqueeze(0)
|
||||
|
||||
return pattern
|
||||
|
||||
def update(
|
||||
self,
|
||||
data_type,
|
||||
data,
|
||||
flow_field,
|
||||
strength=5.0,
|
||||
mode="bilinear",
|
||||
padding_mode="border",
|
||||
align_corners=True,
|
||||
debug_mode="none",
|
||||
always_execute=True,
|
||||
unique_id=None
|
||||
):
|
||||
"""
|
||||
Warp image or latent using optical flow field.
|
||||
|
||||
Args:
|
||||
data_type: Type of data to warp ('IMAGE' or 'LATENT')
|
||||
data: Image tensor [B, H, W, C] or latent dict
|
||||
flow_field: Flow field tensor [B, 2, H, W]
|
||||
strength: Strength of the flow warping effect
|
||||
mode: Interpolation mode
|
||||
padding_mode: Padding mode
|
||||
align_corners: Whether to align corners in sampling grid
|
||||
debug_mode: Visualization mode for debugging
|
||||
always_execute: Whether to always execute (from ControlNodeBase)
|
||||
unique_id: Unique ID for state management (from ControlNodeBase)
|
||||
|
||||
Returns:
|
||||
Warped image or latent, and debug visualization
|
||||
"""
|
||||
# Print debug info
|
||||
print(f"FlowWarpNode - data_type: {data_type}, strength: {strength}")
|
||||
print(f"Flow field shape: {flow_field.shape}, dtype: {flow_field.dtype}, device: {flow_field.device}")
|
||||
if data_type == "IMAGE":
|
||||
print(f"Image shape: {data.shape}, dtype: {data.dtype}, device: {data.device}")
|
||||
|
||||
# Flow statistics
|
||||
with torch.no_grad():
|
||||
flow_abs_mean = torch.abs(flow_field).mean().item()
|
||||
flow_min = flow_field.min().item()
|
||||
flow_max = flow_field.max().item()
|
||||
print(f"Flow stats - mean abs: {flow_abs_mean:.4f}, min: {flow_min:.4f}, max: {flow_max:.4f}")
|
||||
|
||||
# Scale flow by strength
|
||||
scaled_flow = flow_field * strength
|
||||
|
||||
# Create debug visualization based on mode
|
||||
if data_type == "IMAGE":
|
||||
if debug_mode == "flow":
|
||||
# Flow RGB visualization
|
||||
flow_rgb = flow_to_rgb(scaled_flow)
|
||||
debug_vis = flow_rgb.permute(0, 2, 3, 1) # BCHW -> BHWC
|
||||
elif debug_mode == "test_pattern":
|
||||
# Use test pattern to clearly see warping effect
|
||||
B, H, W, C = data.shape
|
||||
test_pattern = self.create_test_pattern(H, W, device=data.device)
|
||||
data = test_pattern # Replace input with test pattern
|
||||
# Also show this as debug visualization
|
||||
debug_vis = test_pattern
|
||||
else:
|
||||
# No specific debug visualization
|
||||
debug_vis = torch.zeros_like(data)
|
||||
else:
|
||||
# For latent, just create a placeholder visualization
|
||||
debug_vis = torch.zeros((1, 128, 128, 3), device=flow_field.device)
|
||||
|
||||
# Process based on data type
|
||||
if data_type == "IMAGE":
|
||||
B, H, W, C = data.shape
|
||||
device = data.device
|
||||
|
||||
# Create coordinate grid in pixel space (0 to H-1, 0 to W-1)
|
||||
y_coords = torch.arange(0, H, device=device).float()
|
||||
x_coords = torch.arange(0, W, device=device).float()
|
||||
grid_y, grid_x = torch.meshgrid(y_coords, x_coords, indexing="ij")
|
||||
pixel_grid = torch.stack([grid_x, grid_y], dim=-1) # [H, W, 2]
|
||||
|
||||
# Apply flow to pixel coordinates (forward warping)
|
||||
# But we need backward warping for grid_sample, so we need to convert
|
||||
# Flow is "where pixels move to" but we need "where to sample from"
|
||||
|
||||
# Convert flow from BCHW [B, 2, H, W] to BHWC [B, H, W, 2]
|
||||
flow_grid = scaled_flow.permute(0, 2, 3, 1).clone() # Clone to avoid in-place issues
|
||||
|
||||
# Apply flow to destination coordinates
|
||||
# dst_coords = pixel_grid + flow_grid
|
||||
|
||||
# For grid_sample, we need normalized coords in [-1, 1]
|
||||
# First, build a basic normalized grid
|
||||
norm_y = 2.0 * torch.arange(H, device=device) / (H - 1) - 1.0
|
||||
norm_x = 2.0 * torch.arange(W, device=device) / (W - 1) - 1.0
|
||||
norm_grid_y, norm_grid_x = torch.meshgrid(norm_y, norm_x, indexing="ij")
|
||||
identity_grid = torch.stack([norm_grid_x, norm_grid_y], dim=-1) # [H, W, 2]
|
||||
|
||||
# Expand to batch dimension
|
||||
identity_grid = identity_grid.expand(B, H, W, 2)
|
||||
|
||||
# For grid_sample, we need to know where to sample from
|
||||
# Flow tells us where pixels move to, so we reverse it
|
||||
# Normalize flow to [-1, 1] range
|
||||
flow_grid[..., 0] = flow_grid[..., 0] * (2.0 / (W - 1)) # x component
|
||||
flow_grid[..., 1] = flow_grid[..., 1] * (2.0 / (H - 1)) # y component
|
||||
|
||||
# Build the sampling grid by subtracting normalized flow from identity grid
|
||||
sampling_grid = identity_grid - flow_grid
|
||||
|
||||
if debug_mode == "grid" and data_type == "IMAGE":
|
||||
# Visualize grid distortion
|
||||
debug_img = torch.zeros((B, H, W, 3), device=device)
|
||||
|
||||
# Normalize grid to [0, 1] for visualization
|
||||
grid_vis = (sampling_grid + 1) / 2.0
|
||||
|
||||
# Assign x and y components to R and G channels
|
||||
debug_img[..., 0] = grid_vis[..., 0] # R = x
|
||||
debug_img[..., 1] = grid_vis[..., 1] # G = y
|
||||
|
||||
# Draw grid lines
|
||||
grid_size = 0.1
|
||||
grid_x = torch.abs(torch.frac(grid_vis[..., 0] / grid_size) - 0.5) < 0.02
|
||||
grid_y = torch.abs(torch.frac(grid_vis[..., 1] / grid_size) - 0.5) < 0.02
|
||||
debug_img[..., 2] = torch.logical_or(grid_x, grid_y).float()
|
||||
|
||||
debug_vis = debug_img
|
||||
|
||||
# Convert image from BHWC to BCHW for grid_sample
|
||||
image_bchw = data.permute(0, 3, 1, 2)
|
||||
|
||||
# Apply grid sampling
|
||||
warped_image_bchw = F.grid_sample(
|
||||
image_bchw,
|
||||
sampling_grid,
|
||||
mode=mode,
|
||||
padding_mode=padding_mode,
|
||||
align_corners=align_corners
|
||||
)
|
||||
|
||||
# Convert back to BHWC
|
||||
warped_image = warped_image_bchw.permute(0, 2, 3, 1)
|
||||
|
||||
# For test pattern, return both the warped pattern and the original
|
||||
if debug_mode == "test_pattern":
|
||||
# Show side-by-side comparison
|
||||
combined = torch.cat([data, warped_image], dim=2) # Concatenate along width
|
||||
debug_vis = combined
|
||||
|
||||
return (warped_image, debug_vis)
|
||||
|
||||
elif data_type == "LATENT":
|
||||
# Handle LATENT data type
|
||||
if not isinstance(data, dict) or "samples" not in data:
|
||||
raise ValueError("Expected latent input to be a dictionary with 'samples' key")
|
||||
|
||||
samples = data["samples"]
|
||||
B, C, H, W = samples.shape
|
||||
device = samples.device
|
||||
|
||||
# Resize flow to match latent dimensions if needed
|
||||
if flow_field.shape[2:] != samples.shape[2:]:
|
||||
scaled_flow = F.interpolate(
|
||||
scaled_flow,
|
||||
size=(H, W),
|
||||
mode='bilinear',
|
||||
align_corners=True
|
||||
)
|
||||
|
||||
# Build the sampling grid
|
||||
norm_y = 2.0 * torch.arange(H, device=device) / (H - 1) - 1.0
|
||||
norm_x = 2.0 * torch.arange(W, device=device) / (W - 1) - 1.0
|
||||
norm_grid_y, norm_grid_x = torch.meshgrid(norm_y, norm_x, indexing="ij")
|
||||
identity_grid = torch.stack([norm_grid_x, norm_grid_y], dim=-1) # [H, W, 2]
|
||||
|
||||
# Expand to batch dimension
|
||||
identity_grid = identity_grid.expand(B, H, W, 2)
|
||||
|
||||
# Convert flow from BCHW [B, 2, H, W] to BHWC [B, H, W, 2]
|
||||
flow_grid = scaled_flow.permute(0, 2, 3, 1).clone()
|
||||
|
||||
# Normalize flow to [-1, 1] range
|
||||
flow_grid[..., 0] = flow_grid[..., 0] * (2.0 / (W - 1)) # x component
|
||||
flow_grid[..., 1] = flow_grid[..., 1] * (2.0 / (H - 1)) # y component
|
||||
|
||||
# Build the sampling grid by subtracting normalized flow from identity grid
|
||||
sampling_grid = identity_grid - flow_grid
|
||||
|
||||
# Apply grid sampling (latent is already in BCHW format)
|
||||
warped_samples = F.grid_sample(
|
||||
samples,
|
||||
sampling_grid,
|
||||
mode=mode,
|
||||
padding_mode=padding_mode,
|
||||
align_corners=align_corners
|
||||
)
|
||||
|
||||
# Create new latent dict with warped samples
|
||||
warped_latent = {k: v for k, v in data.items()} # Shallow copy
|
||||
warped_latent["samples"] = warped_samples
|
||||
|
||||
return (warped_latent, debug_vis)
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unsupported data type: {data_type}")
|
||||
@@ -0,0 +1,327 @@
|
||||
"""
|
||||
Transform control nodes for manipulating and creating transformation matrices.
|
||||
|
||||
These nodes allow for precise control of affine transformations in real-time.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import math
|
||||
|
||||
from ....src.realtimenodes.control_base import ControlNodeBase
|
||||
from ....src.utils.transforms import (
|
||||
create_identity_matrix,
|
||||
create_translation_matrix,
|
||||
create_rotation_matrix,
|
||||
create_scale_matrix,
|
||||
compose_transforms
|
||||
)
|
||||
|
||||
|
||||
class TransformMatrixNode(ControlNodeBase):
|
||||
"""
|
||||
Node that maintains and manipulates an affine transformation matrix with precise control.
|
||||
|
||||
Allows for precise mathematical operations on transformation matrices.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
inputs = super().INPUT_TYPES()
|
||||
inputs["required"].update({
|
||||
"operation": (
|
||||
["initialize", "translate", "rotate", "scale", "compose", "reset"],
|
||||
{"default": "initialize"}
|
||||
),
|
||||
"translate_x": (
|
||||
"FLOAT",
|
||||
{"default": 0.0, "step": 0.01, "tooltip": "Translation in x direction (-1 to 1 range)"}
|
||||
),
|
||||
"translate_y": (
|
||||
"FLOAT",
|
||||
{"default": 0.0, "step": 0.01, "tooltip": "Translation in y direction (-1 to 1 range)"}
|
||||
),
|
||||
"rotate_angle": (
|
||||
"FLOAT",
|
||||
{"default": 0.0, "step": 0.1, "tooltip": "Rotation angle in degrees"}
|
||||
),
|
||||
"scale_factor_x": (
|
||||
"FLOAT",
|
||||
{"default": 1.0, "step": 0.01, "tooltip": "Scale factor in x direction"}
|
||||
),
|
||||
"scale_factor_y": (
|
||||
"FLOAT",
|
||||
{"default": 1.0, "step": 0.01, "tooltip": "Scale factor in y direction"}
|
||||
),
|
||||
"center_x": (
|
||||
"FLOAT",
|
||||
{"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01,
|
||||
"tooltip": "X coordinate of transformation center (0-1 range)"}
|
||||
),
|
||||
"center_y": (
|
||||
"FLOAT",
|
||||
{"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01,
|
||||
"tooltip": "Y coordinate of transformation center (0-1 range)"}
|
||||
),
|
||||
})
|
||||
inputs["optional"] = {
|
||||
"secondary_matrix": (
|
||||
"TRANSFORM_MATRIX",
|
||||
{"tooltip": "Secondary matrix for composition operations"}
|
||||
)
|
||||
}
|
||||
return inputs
|
||||
|
||||
RETURN_TYPES = ("TRANSFORM_MATRIX",)
|
||||
RETURN_NAMES = ("transform_matrix",)
|
||||
FUNCTION = "update"
|
||||
CATEGORY = "Realtime Nodes/control/transform"
|
||||
|
||||
def update(
|
||||
self,
|
||||
operation,
|
||||
translate_x,
|
||||
translate_y,
|
||||
rotate_angle,
|
||||
scale_factor_x,
|
||||
scale_factor_y,
|
||||
center_x,
|
||||
center_y,
|
||||
secondary_matrix=None,
|
||||
always_execute=True,
|
||||
unique_id=None
|
||||
):
|
||||
"""
|
||||
Update the transformation matrix based on the specified operation.
|
||||
|
||||
Args:
|
||||
operation: The operation to perform on the matrix
|
||||
translate_x: X translation amount
|
||||
translate_y: Y translation amount
|
||||
rotate_angle: Rotation angle in degrees
|
||||
scale_factor_x: X scale factor
|
||||
scale_factor_y: Y scale factor
|
||||
center_x: X coordinate of transformation center
|
||||
center_y: Y coordinate of transformation center
|
||||
secondary_matrix: Optional secondary matrix for composition
|
||||
always_execute: Whether to always execute (from ControlNodeBase)
|
||||
unique_id: Unique ID for state management (from ControlNodeBase)
|
||||
|
||||
Returns:
|
||||
Updated transformation matrix
|
||||
"""
|
||||
# Get state or create identity matrix
|
||||
state = self.get_state({"matrix": None}, unique_id)
|
||||
|
||||
# Default to CPU if we don't have a device yet
|
||||
device = "cpu"
|
||||
if state["matrix"] is not None:
|
||||
device = state["matrix"].device
|
||||
elif secondary_matrix is not None:
|
||||
device = secondary_matrix.device
|
||||
|
||||
current_matrix = state["matrix"] if state["matrix"] is not None else create_identity_matrix(device)
|
||||
|
||||
# Perform operation
|
||||
if operation == "initialize":
|
||||
current_matrix = create_identity_matrix(device)
|
||||
|
||||
elif operation == "translate":
|
||||
t_matrix = create_translation_matrix(translate_x, translate_y, device)
|
||||
current_matrix = compose_transforms(t_matrix, current_matrix)
|
||||
|
||||
elif operation == "rotate":
|
||||
r_matrix = create_rotation_matrix(rotate_angle, center_x, center_y, device)
|
||||
current_matrix = compose_transforms(r_matrix, current_matrix)
|
||||
|
||||
elif operation == "scale":
|
||||
s_matrix = create_scale_matrix(scale_factor_x, scale_factor_y, center_x, center_y, device)
|
||||
current_matrix = compose_transforms(s_matrix, current_matrix)
|
||||
|
||||
elif operation == "compose" and secondary_matrix is not None:
|
||||
current_matrix = compose_transforms(secondary_matrix, current_matrix)
|
||||
|
||||
elif operation == "reset":
|
||||
current_matrix = create_identity_matrix(device)
|
||||
|
||||
# Update state
|
||||
state["matrix"] = current_matrix
|
||||
self.set_state(state, unique_id)
|
||||
|
||||
return (current_matrix,)
|
||||
|
||||
|
||||
class CameraControlNode(ControlNodeBase):
|
||||
"""
|
||||
High-level camera control with intuitive, independent parameters.
|
||||
|
||||
Provides artist-friendly controls for camera manipulation in real-time.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
inputs = super().INPUT_TYPES()
|
||||
inputs["required"].update({
|
||||
"pan_x": (
|
||||
"FLOAT",
|
||||
{"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01,
|
||||
"tooltip": "Horizontal panning (-1.0 to 1.0)"}
|
||||
),
|
||||
"pan_y": (
|
||||
"FLOAT",
|
||||
{"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01,
|
||||
"tooltip": "Vertical panning (-1.0 to 1.0)"}
|
||||
),
|
||||
"zoom": (
|
||||
"FLOAT",
|
||||
{"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.01,
|
||||
"tooltip": "Zoom factor (1.0 = no zoom)"}
|
||||
),
|
||||
"rotation": (
|
||||
"FLOAT",
|
||||
{"default": 0.0, "min": -180.0, "max": 180.0, "step": 0.1,
|
||||
"tooltip": "Rotation in degrees"}
|
||||
),
|
||||
"center_x": (
|
||||
"FLOAT",
|
||||
{"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01,
|
||||
"tooltip": "X coordinate of rotation/zoom center"}
|
||||
),
|
||||
"center_y": (
|
||||
"FLOAT",
|
||||
{"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01,
|
||||
"tooltip": "Y coordinate of rotation/zoom center"}
|
||||
),
|
||||
"reset": (
|
||||
"BOOLEAN",
|
||||
{"default": False, "tooltip": "Reset camera to default state"}
|
||||
),
|
||||
"accumulate": (
|
||||
"BOOLEAN",
|
||||
{"default": False, "tooltip": "Accumulate transformations over time (vs. absolute values)"}
|
||||
),
|
||||
})
|
||||
inputs["optional"] = {
|
||||
"base_matrix": (
|
||||
"TRANSFORM_MATRIX",
|
||||
{"tooltip": "Starting transform matrix (optional)"}
|
||||
)
|
||||
}
|
||||
return inputs
|
||||
|
||||
RETURN_TYPES = ("TRANSFORM_MATRIX",)
|
||||
RETURN_NAMES = ("transform_matrix",)
|
||||
FUNCTION = "update"
|
||||
CATEGORY = "Realtime Nodes/control/transform"
|
||||
|
||||
def update(
|
||||
self,
|
||||
pan_x,
|
||||
pan_y,
|
||||
zoom,
|
||||
rotation,
|
||||
center_x,
|
||||
center_y,
|
||||
reset,
|
||||
accumulate,
|
||||
base_matrix=None,
|
||||
always_execute=True,
|
||||
unique_id=None
|
||||
):
|
||||
"""
|
||||
Update the camera transformation matrix.
|
||||
|
||||
Args:
|
||||
pan_x: Horizontal panning amount
|
||||
pan_y: Vertical panning amount
|
||||
zoom: Zoom factor
|
||||
rotation: Rotation angle in degrees
|
||||
center_x: X coordinate of rotation/zoom center
|
||||
center_y: Y coordinate of rotation/zoom center
|
||||
reset: Whether to reset the camera
|
||||
accumulate: Whether to accumulate transforms (vs. absolute values)
|
||||
base_matrix: Optional base matrix to start from
|
||||
always_execute: Whether to always execute (from ControlNodeBase)
|
||||
unique_id: Unique ID for state management (from ControlNodeBase)
|
||||
|
||||
Returns:
|
||||
Updated camera transformation matrix
|
||||
"""
|
||||
# Get state or initialize
|
||||
state = self.get_state({
|
||||
"matrix": None, # Current transform matrix
|
||||
"last_params": { # Last used parameters
|
||||
"pan_x": 0.0,
|
||||
"pan_y": 0.0,
|
||||
"zoom": 1.0,
|
||||
"rotation": 0.0
|
||||
}
|
||||
}, unique_id)
|
||||
|
||||
device = "cpu"
|
||||
|
||||
# Handle reset
|
||||
if reset:
|
||||
state["matrix"] = None
|
||||
state["last_params"] = {
|
||||
"pan_x": 0.0,
|
||||
"pan_y": 0.0,
|
||||
"zoom": 1.0,
|
||||
"rotation": 0.0
|
||||
}
|
||||
|
||||
# Start with base_matrix or identity
|
||||
if base_matrix is not None:
|
||||
current_matrix = base_matrix
|
||||
device = base_matrix.device
|
||||
elif not accumulate:
|
||||
# For non-accumulate mode, always start fresh
|
||||
current_matrix = create_identity_matrix(device)
|
||||
elif state["matrix"] is not None:
|
||||
current_matrix = state["matrix"]
|
||||
device = current_matrix.device
|
||||
else:
|
||||
current_matrix = create_identity_matrix(device)
|
||||
|
||||
# For accumulate mode, calculate the delta from last parameters
|
||||
if accumulate and state["last_params"] is not None:
|
||||
# Calculate deltas
|
||||
delta_pan_x = pan_x - state["last_params"]["pan_x"]
|
||||
delta_pan_y = pan_y - state["last_params"]["pan_y"]
|
||||
delta_zoom = zoom / state["last_params"]["zoom"] if state["last_params"]["zoom"] != 0 else zoom
|
||||
delta_rotation = rotation - state["last_params"]["rotation"]
|
||||
|
||||
# Use deltas instead of absolute values
|
||||
pan_x = delta_pan_x
|
||||
pan_y = delta_pan_y
|
||||
zoom = delta_zoom
|
||||
rotation = delta_rotation
|
||||
|
||||
# Apply transforms in sequence: first scale (zoom), then rotate, then translate (pan)
|
||||
# This order ensures consistent behavior
|
||||
|
||||
# 1. Apply zoom
|
||||
if zoom != 1.0:
|
||||
zoom_matrix = create_scale_matrix(zoom, zoom, center_x, center_y, device)
|
||||
current_matrix = compose_transforms(zoom_matrix, current_matrix)
|
||||
|
||||
# 2. Apply rotation
|
||||
if rotation != 0.0:
|
||||
rotation_matrix = create_rotation_matrix(rotation, center_x, center_y, device)
|
||||
current_matrix = compose_transforms(rotation_matrix, current_matrix)
|
||||
|
||||
# 3. Apply panning
|
||||
if pan_x != 0.0 or pan_y != 0.0:
|
||||
translation_matrix = create_translation_matrix(pan_x, pan_y, device)
|
||||
current_matrix = compose_transforms(translation_matrix, current_matrix)
|
||||
|
||||
# Update state with current parameters and matrix
|
||||
state["matrix"] = current_matrix
|
||||
state["last_params"] = {
|
||||
"pan_x": pan_x if not accumulate else pan_x + state["last_params"]["pan_x"],
|
||||
"pan_y": pan_y if not accumulate else pan_y + state["last_params"]["pan_y"],
|
||||
"zoom": zoom if not accumulate else zoom * state["last_params"]["zoom"],
|
||||
"rotation": rotation if not accumulate else rotation + state["last_params"]["rotation"]
|
||||
}
|
||||
self.set_state(state, unique_id)
|
||||
|
||||
return (current_matrix,)
|
||||
@@ -0,0 +1,152 @@
|
||||
"""
|
||||
Transform warping nodes for applying transformation matrices to images and latents.
|
||||
|
||||
These nodes apply affine transformations to various media types in real-time.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from ....src.realtimenodes.control_base import ControlNodeBase
|
||||
from ....src.utils.warp import warp_image, warp_latent
|
||||
|
||||
#TODO add depth aware throughout
|
||||
|
||||
class MatrixTransformImageNode(ControlNodeBase):
|
||||
"""
|
||||
Apply affine transformation matrix to images.
|
||||
|
||||
Warps images using a provided transformation matrix.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
inputs = super().INPUT_TYPES()
|
||||
inputs["required"].update({
|
||||
"image": ("IMAGE",),
|
||||
"transform_matrix": ("TRANSFORM_MATRIX",),
|
||||
"mode": (
|
||||
["bilinear", "nearest", "bicubic"],
|
||||
{"default": "bilinear", "tooltip": "Interpolation mode for sampling"}
|
||||
),
|
||||
"padding_mode": (
|
||||
["zeros", "border", "reflection"],
|
||||
{"default": "border", "tooltip": "How to handle pixels outside the image boundaries"}
|
||||
),
|
||||
"align_corners": (
|
||||
"BOOLEAN",
|
||||
{"default": True, "tooltip": "Whether to align corners when sampling"}
|
||||
),
|
||||
})
|
||||
return inputs
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("transformed_image",)
|
||||
FUNCTION = "update"
|
||||
CATEGORY = "Realtime Nodes/control/transform"
|
||||
|
||||
def update(
|
||||
self,
|
||||
image,
|
||||
transform_matrix,
|
||||
mode="bilinear",
|
||||
padding_mode="border",
|
||||
align_corners=True,
|
||||
always_execute=True,
|
||||
unique_id=None
|
||||
):
|
||||
"""
|
||||
Apply transformation matrix to an image.
|
||||
|
||||
Args:
|
||||
image: Input image tensor [B, H, W, C]
|
||||
transform_matrix: 3x3 transformation matrix
|
||||
mode: Interpolation mode
|
||||
padding_mode: Padding mode
|
||||
align_corners: Whether to align corners in sampling
|
||||
always_execute: Whether to always execute (from ControlNodeBase)
|
||||
unique_id: Unique ID for state management (from ControlNodeBase)
|
||||
|
||||
Returns:
|
||||
Warped image
|
||||
"""
|
||||
# Apply warping to image using utility function
|
||||
warped_image = warp_image(
|
||||
image,
|
||||
transform_matrix,
|
||||
mode=mode,
|
||||
padding_mode=padding_mode,
|
||||
align_corners=align_corners
|
||||
)
|
||||
|
||||
return (warped_image,)
|
||||
|
||||
|
||||
class MatrixTransformLatentNode(ControlNodeBase):
|
||||
"""
|
||||
Apply affine transformation matrix to latent space.
|
||||
|
||||
Warps latent tensors using a provided transformation matrix.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
inputs = super().INPUT_TYPES()
|
||||
inputs["required"].update({
|
||||
"latent": ("LATENT",),
|
||||
"transform_matrix": ("TRANSFORM_MATRIX",),
|
||||
"mode": (
|
||||
["bilinear", "nearest", "bicubic"],
|
||||
{"default": "bilinear", "tooltip": "Interpolation mode for sampling"}
|
||||
),
|
||||
"padding_mode": (
|
||||
["zeros", "border", "reflection"],
|
||||
{"default": "border", "tooltip": "How to handle values outside the latent boundaries"}
|
||||
),
|
||||
"align_corners": (
|
||||
"BOOLEAN",
|
||||
{"default": True, "tooltip": "Whether to align corners when sampling"}
|
||||
),
|
||||
})
|
||||
return inputs
|
||||
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
RETURN_NAMES = ("transformed_latent",)
|
||||
FUNCTION = "update"
|
||||
CATEGORY = "Realtime Nodes/control/transform"
|
||||
|
||||
def update(
|
||||
self,
|
||||
latent,
|
||||
transform_matrix,
|
||||
mode="bilinear",
|
||||
padding_mode="border",
|
||||
align_corners=True,
|
||||
always_execute=True,
|
||||
unique_id=None
|
||||
):
|
||||
"""
|
||||
Apply transformation matrix to latent space.
|
||||
|
||||
Args:
|
||||
latent: Input latent dict with 'samples' tensor [B, C, H, W]
|
||||
transform_matrix: 3x3 transformation matrix
|
||||
mode: Interpolation mode
|
||||
padding_mode: Padding mode
|
||||
align_corners: Whether to align corners in sampling
|
||||
always_execute: Whether to always execute (from ControlNodeBase)
|
||||
unique_id: Unique ID for state management (from ControlNodeBase)
|
||||
|
||||
Returns:
|
||||
Warped latent
|
||||
"""
|
||||
# Apply warping to latent using utility function
|
||||
warped_latent = warp_latent(
|
||||
latent,
|
||||
transform_matrix,
|
||||
mode=mode,
|
||||
padding_mode=padding_mode,
|
||||
align_corners=align_corners
|
||||
)
|
||||
|
||||
return (warped_latent,)
|
||||
+26
-1
@@ -6,7 +6,7 @@ Unified utilities package for ComfyUI RealTimeNodes.
|
||||
from .general import AlwaysEqualProxy
|
||||
|
||||
# Image utilities
|
||||
from .image import convert_to_cv2, convert_to_tensor, create_mask_from_points
|
||||
from .image import convert_to_cv2, convert_to_tensor, create_mask_from_points, flow_to_rgb, gaussian_blur_2d
|
||||
|
||||
# Math utilities
|
||||
from .math import (
|
||||
@@ -18,6 +18,31 @@ from .math import (
|
||||
# Timing utilities
|
||||
from .timing import TimestampProvider
|
||||
|
||||
# Similar image filter
|
||||
from .similar_image_filter import SimilarImageFilter
|
||||
|
||||
# Transform utilities
|
||||
from .transforms import (
|
||||
create_identity_matrix,
|
||||
create_translation_matrix,
|
||||
create_rotation_matrix,
|
||||
create_scale_matrix,
|
||||
compose_transforms,
|
||||
normalize_homogeneous_coordinates,
|
||||
apply_transform
|
||||
)
|
||||
|
||||
# Warp utilities
|
||||
from .warp import (
|
||||
create_sampling_grid,
|
||||
transform_grid,
|
||||
warp_image,
|
||||
warp_latent
|
||||
)
|
||||
|
||||
# Real-time flownets
|
||||
from .realtime_flownets import RealTimeFlowNet
|
||||
|
||||
# For backward compatibility, expose the original names
|
||||
# This allows existing code to continue working while we transition
|
||||
image_to_cv2 = convert_to_cv2
|
||||
|
||||
@@ -0,0 +1,174 @@
|
||||
"""
|
||||
Transform utilities for manipulating affine transformation matrices.
|
||||
|
||||
Contains functions for creating and composing transformation matrices.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import math
|
||||
|
||||
|
||||
def create_identity_matrix(device="cpu"):
|
||||
"""
|
||||
Create a 3x3 identity transformation matrix.
|
||||
|
||||
Args:
|
||||
device: The torch device to place the matrix on
|
||||
|
||||
Returns:
|
||||
3x3 identity matrix as torch.Tensor
|
||||
"""
|
||||
return torch.eye(3, device=device)
|
||||
|
||||
|
||||
def create_translation_matrix(dx, dy, device="cpu"):
|
||||
"""
|
||||
Create a 3x3 translation matrix.
|
||||
|
||||
Args:
|
||||
dx: Translation in x direction
|
||||
dy: Translation in y direction
|
||||
device: The torch device to place the matrix on
|
||||
|
||||
Returns:
|
||||
3x3 translation matrix as torch.Tensor
|
||||
"""
|
||||
matrix = torch.eye(3, device=device)
|
||||
matrix[0, 2] = dx
|
||||
matrix[1, 2] = dy
|
||||
return matrix
|
||||
|
||||
|
||||
def create_rotation_matrix(angle_degrees, center_x=0.5, center_y=0.5, device="cpu"):
|
||||
"""
|
||||
Create a 3x3 rotation matrix that rotates around the specified center point.
|
||||
|
||||
Args:
|
||||
angle_degrees: Rotation angle in degrees
|
||||
center_x: X coordinate of rotation center (0-1 range)
|
||||
center_y: Y coordinate of rotation center (0-1 range)
|
||||
device: The torch device to place the matrix on
|
||||
|
||||
Returns:
|
||||
3x3 rotation matrix as torch.Tensor
|
||||
"""
|
||||
angle_rad = math.radians(angle_degrees)
|
||||
cos_theta = math.cos(angle_rad)
|
||||
sin_theta = math.sin(angle_rad)
|
||||
|
||||
# Convert center from 0-1 range to -1 to 1 range
|
||||
center_x = (center_x * 2) - 1
|
||||
center_y = (center_y * 2) - 1
|
||||
|
||||
# 1. Translate center to origin
|
||||
t1 = create_translation_matrix(-center_x, -center_y, device)
|
||||
|
||||
# 2. Rotate around origin
|
||||
r = torch.eye(3, device=device)
|
||||
r[0, 0] = cos_theta
|
||||
r[0, 1] = -sin_theta
|
||||
r[1, 0] = sin_theta
|
||||
r[1, 1] = cos_theta
|
||||
|
||||
# 3. Translate back
|
||||
t2 = create_translation_matrix(center_x, center_y, device)
|
||||
|
||||
# Compose transformations: t2 @ r @ t1
|
||||
return torch.matmul(t2, torch.matmul(r, t1))
|
||||
|
||||
|
||||
def create_scale_matrix(scale_x, scale_y, center_x=0.5, center_y=0.5, device="cpu"):
|
||||
"""
|
||||
Create a 3x3 scaling matrix that scales around the specified center point.
|
||||
|
||||
Args:
|
||||
scale_x: Scale factor in x direction
|
||||
scale_y: Scale factor in y direction
|
||||
center_x: X coordinate of scaling center (0-1 range)
|
||||
center_y: Y coordinate of scaling center (0-1 range)
|
||||
device: The torch device to place the matrix on
|
||||
|
||||
Returns:
|
||||
3x3 scaling matrix as torch.Tensor
|
||||
"""
|
||||
# Convert center from 0-1 range to -1 to 1 range
|
||||
center_x = (center_x * 2) - 1
|
||||
center_y = (center_y * 2) - 1
|
||||
|
||||
# 1. Translate center to origin
|
||||
t1 = create_translation_matrix(-center_x, -center_y, device)
|
||||
|
||||
# 2. Scale around origin
|
||||
s = torch.eye(3, device=device)
|
||||
s[0, 0] = scale_x
|
||||
s[1, 1] = scale_y
|
||||
|
||||
# 3. Translate back
|
||||
t2 = create_translation_matrix(center_x, center_y, device)
|
||||
|
||||
# Compose transformations: t2 @ s @ t1
|
||||
return torch.matmul(t2, torch.matmul(s, t1))
|
||||
|
||||
|
||||
def compose_transforms(transform_a, transform_b):
|
||||
"""
|
||||
Compose two transformation matrices. The result is equivalent to applying
|
||||
transform_b first, then transform_a.
|
||||
|
||||
Args:
|
||||
transform_a: First transformation matrix
|
||||
transform_b: Second transformation matrix
|
||||
|
||||
Returns:
|
||||
Composed transformation matrix
|
||||
"""
|
||||
return torch.matmul(transform_a, transform_b)
|
||||
|
||||
|
||||
def normalize_homogeneous_coordinates(coords):
|
||||
"""
|
||||
Normalize homogeneous coordinates by dividing by the last component.
|
||||
|
||||
Args:
|
||||
coords: Tensor of homogeneous coordinates [..., 3]
|
||||
|
||||
Returns:
|
||||
Normalized coordinates
|
||||
"""
|
||||
# Extract last component
|
||||
w = coords[..., 2:3]
|
||||
|
||||
# Handle zeros (avoid division by zero)
|
||||
mask = w != 0
|
||||
|
||||
# Create a copy of the coordinates
|
||||
normalized = coords.clone()
|
||||
|
||||
# Apply normalization only to non-zero denominator elements
|
||||
# We need to expand the mask to match the coords dimension
|
||||
for i in range(coords.shape[-1]):
|
||||
normalized[..., i] = torch.where(mask.squeeze(-1), coords[..., i] / w.squeeze(-1), coords[..., i])
|
||||
|
||||
return normalized
|
||||
|
||||
|
||||
def apply_transform(points, transform_matrix):
|
||||
"""
|
||||
Apply a transformation matrix to a set of points.
|
||||
|
||||
Args:
|
||||
points: Tensor of points [..., 2]
|
||||
transform_matrix: 3x3 transformation matrix
|
||||
|
||||
Returns:
|
||||
Transformed points
|
||||
"""
|
||||
# Convert to homogeneous coordinates
|
||||
homogeneous = torch.cat([points, torch.ones_like(points[..., :1])], dim=-1)
|
||||
|
||||
# Apply transformation
|
||||
transformed = torch.matmul(homogeneous, transform_matrix.T)
|
||||
|
||||
# Normalize and return x, y coordinates
|
||||
normalized = normalize_homogeneous_coordinates(transformed)
|
||||
return normalized[..., :2]
|
||||
@@ -0,0 +1,181 @@
|
||||
"""
|
||||
Warping utilities for transforming images and latents using transformation matrices.
|
||||
|
||||
Contains functions for grid-based warping operations.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from .transforms import normalize_homogeneous_coordinates
|
||||
|
||||
|
||||
def create_sampling_grid(height, width, device="cpu"):
|
||||
"""
|
||||
Create a normalized sampling grid for an image of the given dimensions.
|
||||
|
||||
Args:
|
||||
height: Image height
|
||||
width: Image width
|
||||
device: The torch device to place the grid on
|
||||
|
||||
Returns:
|
||||
Normalized grid of shape [height, width, 2] with values in range [-1, 1]
|
||||
"""
|
||||
# Create coordinate grid in pixel space
|
||||
y_coords = torch.linspace(0, height - 1, height, device=device)
|
||||
x_coords = torch.linspace(0, width - 1, width, device=device)
|
||||
grid_y, grid_x = torch.meshgrid(y_coords, x_coords, indexing="ij")
|
||||
|
||||
# Convert to normalized coordinates in range [-1, 1]
|
||||
grid_x = 2.0 * grid_x / (width - 1) - 1.0
|
||||
grid_y = 2.0 * grid_y / (height - 1) - 1.0
|
||||
|
||||
# Stack into grid of shape [height, width, 2]
|
||||
grid = torch.stack([grid_x, grid_y], dim=-1)
|
||||
|
||||
return grid
|
||||
|
||||
|
||||
def transform_grid(grid, transform_matrix):
|
||||
"""
|
||||
Apply a transformation matrix to a sampling grid.
|
||||
|
||||
The transformation matrix is applied in reverse (inverse) since we're creating
|
||||
a sampling grid - we need to know where to sample from in the source image.
|
||||
|
||||
Args:
|
||||
grid: Sampling grid of shape [..., 2]
|
||||
transform_matrix: 3x3 transformation matrix
|
||||
|
||||
Returns:
|
||||
Transformed grid of same shape as input
|
||||
"""
|
||||
# Get grid shape for reshaping later
|
||||
original_shape = grid.shape
|
||||
device = grid.device
|
||||
|
||||
# Flatten the grid to make computations easier
|
||||
# Reshape to [N, 2] where N is the product of all dimensions except the last
|
||||
flat_shape = (-1, 2)
|
||||
grid_flat = grid.reshape(flat_shape)
|
||||
|
||||
# Convert to homogeneous coordinates
|
||||
ones = torch.ones(grid_flat.shape[0], 1, device=device)
|
||||
homogeneous_grid = torch.cat([grid_flat, ones], dim=-1)
|
||||
|
||||
# Compute inverse for sampling (where to get pixels from)
|
||||
try:
|
||||
inverse_matrix = torch.inverse(transform_matrix)
|
||||
except RuntimeError:
|
||||
# Handle singular matrix by adding small epsilon
|
||||
epsilon = 1e-6
|
||||
identity = torch.eye(3, device=device) * epsilon
|
||||
transform_matrix_reg = transform_matrix + identity
|
||||
inverse_matrix = torch.inverse(transform_matrix_reg)
|
||||
|
||||
# Apply transformation
|
||||
transformed = torch.matmul(homogeneous_grid, inverse_matrix.T)
|
||||
|
||||
# Normalize homogeneous coordinates (divide by w component)
|
||||
w = transformed[:, 2:3]
|
||||
mask = w != 0
|
||||
normalized = torch.zeros_like(transformed)
|
||||
|
||||
# Apply normalization to x, y, z components where w != 0
|
||||
normalized[:, 0] = torch.where(mask.squeeze(-1), transformed[:, 0] / w.squeeze(-1), transformed[:, 0])
|
||||
normalized[:, 1] = torch.where(mask.squeeze(-1), transformed[:, 1] / w.squeeze(-1), transformed[:, 1])
|
||||
normalized[:, 2] = torch.where(mask.squeeze(-1), transformed[:, 2] / w.squeeze(-1), transformed[:, 2])
|
||||
|
||||
# Get x, y components and reshape back to original grid shape
|
||||
transformed_grid = normalized[:, :2].reshape(original_shape)
|
||||
|
||||
return transformed_grid
|
||||
|
||||
|
||||
def warp_image(image, transform_matrix, mode="bilinear", padding_mode="zeros", align_corners=True):
|
||||
"""
|
||||
Warp an image using a transformation matrix.
|
||||
|
||||
Args:
|
||||
image: Input image tensor of shape [B, H, W, C]
|
||||
transform_matrix: 3x3 transformation matrix
|
||||
mode: Interpolation mode ('bilinear', 'nearest', or 'bicubic')
|
||||
padding_mode: Padding mode ('zeros', 'border', or 'reflection')
|
||||
align_corners: Whether to align corners in the grid_sample operation
|
||||
|
||||
Returns:
|
||||
Warped image of same shape as input
|
||||
"""
|
||||
# Get image dimensions
|
||||
B, H, W, C = image.shape
|
||||
device = image.device
|
||||
|
||||
# Create sampling grid
|
||||
grid = create_sampling_grid(H, W, device)
|
||||
|
||||
# Apply transformation to grid
|
||||
warped_grid = transform_grid(grid, transform_matrix)
|
||||
|
||||
# Expand grid to batch dimension
|
||||
warped_grid = warped_grid.expand(B, -1, -1, -1)
|
||||
|
||||
# Reorder image for grid_sample: [B, C, H, W]
|
||||
image_bchw = image.permute(0, 3, 1, 2)
|
||||
|
||||
# Apply grid sampling
|
||||
warped_image = F.grid_sample(
|
||||
image_bchw,
|
||||
warped_grid,
|
||||
mode=mode,
|
||||
padding_mode=padding_mode,
|
||||
align_corners=align_corners
|
||||
)
|
||||
|
||||
# Reorder back to [B, H, W, C]
|
||||
warped_image = warped_image.permute(0, 2, 3, 1)
|
||||
|
||||
return warped_image
|
||||
|
||||
|
||||
def warp_latent(latent, transform_matrix, mode="bilinear", padding_mode="zeros", align_corners=True):
|
||||
"""
|
||||
Warp a latent tensor using a transformation matrix.
|
||||
|
||||
Args:
|
||||
latent: Input latent dict with 'samples' tensor of shape [B, C, H, W]
|
||||
transform_matrix: 3x3 transformation matrix
|
||||
mode: Interpolation mode ('bilinear', 'nearest', or 'bicubic')
|
||||
padding_mode: Padding mode ('zeros', 'border', or 'reflection')
|
||||
align_corners: Whether to align corners in the grid_sample operation
|
||||
|
||||
Returns:
|
||||
Warped latent in same format as input
|
||||
"""
|
||||
# Get latent samples
|
||||
samples = latent["samples"]
|
||||
B, C, H, W = samples.shape
|
||||
device = samples.device
|
||||
|
||||
# Create sampling grid
|
||||
grid = create_sampling_grid(H, W, device)
|
||||
|
||||
# Apply transformation to grid
|
||||
warped_grid = transform_grid(grid, transform_matrix)
|
||||
|
||||
# Expand grid to batch dimension
|
||||
warped_grid = warped_grid.expand(B, -1, -1, -1)
|
||||
|
||||
# Apply grid sampling (latent is already in [B, C, H, W] format)
|
||||
warped_samples = F.grid_sample(
|
||||
samples,
|
||||
warped_grid,
|
||||
mode=mode,
|
||||
padding_mode=padding_mode,
|
||||
align_corners=align_corners
|
||||
)
|
||||
|
||||
# Create new latent dict with warped samples
|
||||
warped_latent = latent.copy()
|
||||
warped_latent["samples"] = warped_samples
|
||||
|
||||
return warped_latent
|
||||
Reference in New Issue
Block a user