← back to Handbag Authentication
handbag_data/github_datasets/amazon-berkeley-objects/render/utils/node.py
201 lines
import bpy
import sys
import math
from typing import Iterable
def create_frame_node(node_tree: bpy.types.NodeTree,
nodes: Iterable[bpy.types.Node] = [],
name: str = "Frame",
label: str = "Frame") -> bpy.types.Node:
frame_node = node_tree.nodes.new(type='NodeFrame')
frame_node.name = name
frame_node.label = label
for node in nodes:
node.parent = frame_node
return frame_node
def set_socket_value_range(socket: bpy.types.NodeSocket,
default_value: float = 0.0,
min_value: float = 0.0,
max_value: float = 1.0) -> None:
assert socket.type == "VALUE"
socket.default_value = default_value
socket.min_value = min_value
socket.max_value = max_value
def clean_nodes(nodes: bpy.types.Nodes) -> None:
for node in nodes:
nodes.remove(node)
def arrange_nodes(node_tree: bpy.types.NodeTree, verbose: bool = False) -> None:
max_num_iters = 2000
epsilon = 1e-05
target_space = 50.0
second_stage = False
fix_horizontal_location = True
fix_vertical_location = True
fix_overlaps = True
if verbose:
print("-----------------")
print("Target nodes:")
for node in node_tree.nodes:
print("- " + node.name)
# In the first stage, expand nodes overly
target_space *= 2.0
# Gauss-Seidel-style iterations
previous_squared_deltas_sum = sys.float_info.max
for i in range(max_num_iters):
squared_deltas_sum = 0.0
if fix_horizontal_location:
for link in node_tree.links:
k = 0.9 if not second_stage else 0.5
threshold_factor = 2.0
x_from = link.from_node.location[0]
x_to = link.to_node.location[0]
w_from = link.from_node.width
signed_space = x_to - x_from - w_from
C = signed_space - target_space
grad_C_x_from = -1.0
grad_C_x_to = 1.0
# Skip if the distance is sufficiently large
if C >= target_space * threshold_factor:
continue
lagrange = C / (grad_C_x_from * grad_C_x_from + grad_C_x_to * grad_C_x_to)
delta_x_from = -lagrange * grad_C_x_from
delta_x_to = -lagrange * grad_C_x_to
link.from_node.location[0] += k * delta_x_from
link.to_node.location[0] += k * delta_x_to
squared_deltas_sum += k * k * (delta_x_from * delta_x_from + delta_x_to * delta_x_to)
if fix_vertical_location:
k = 0.5 if not second_stage else 0.05
socket_offset = 20.0
def get_from_socket_index(node: bpy.types.Node, node_socket: bpy.types.NodeSocket) -> int:
for i in range(len(node.outputs)):
if node.outputs[i] == node_socket:
return i
assert False
def get_to_socket_index(node: bpy.types.Node, node_socket: bpy.types.NodeSocket) -> int:
for i in range(len(node.inputs)):
if node.inputs[i] == node_socket:
return i
assert False
for link in node_tree.links:
from_socket_index = get_from_socket_index(link.from_node, link.from_socket)
to_socket_index = get_to_socket_index(link.to_node, link.to_socket)
y_from = link.from_node.location[1] - socket_offset * from_socket_index
y_to = link.to_node.location[1] - socket_offset * to_socket_index
C = y_from - y_to
grad_C_y_from = 1.0
grad_C_y_to = -1.0
lagrange = C / (grad_C_y_from * grad_C_y_from + grad_C_y_to * grad_C_y_to)
delta_y_from = -lagrange * grad_C_y_from
delta_y_to = -lagrange * grad_C_y_to
link.from_node.location[1] += k * delta_y_from
link.to_node.location[1] += k * delta_y_to
squared_deltas_sum += k * k * (delta_y_from * delta_y_from + delta_y_to * delta_y_to)
if fix_overlaps and second_stage:
k = 0.9
margin = 0.5 * target_space
# Examine all node pairs
for node_1 in node_tree.nodes:
for node_2 in node_tree.nodes:
if node_1 == node_2:
continue
x_1 = node_1.location[0]
x_2 = node_2.location[0]
w_1 = node_1.width
w_2 = node_2.width
cx_1 = x_1 + 0.5 * w_1
cx_2 = x_2 + 0.5 * w_2
rx_1 = 0.5 * w_1 + margin
rx_2 = 0.5 * w_2 + margin
# Note: "dimensions" and "height" may not be correct depending on the situation
def get_height(node: bpy.types.Node) -> float:
if node.dimensions.y > epsilon:
return node.dimensions.y
elif math.fabs(node.height - 100.0) > epsilon:
return node.height
else:
return 200.0
y_1 = node_1.location[1]
y_2 = node_2.location[1]
h_1 = get_height(node_1)
h_2 = get_height(node_2)
cy_1 = y_1 - 0.5 * h_1
cy_2 = y_2 - 0.5 * h_2
ry_1 = 0.5 * h_1 + margin
ry_2 = 0.5 * h_2 + margin
C_x = math.fabs(cx_1 - cx_2) - (rx_1 + rx_2)
C_y = math.fabs(cy_1 - cy_2) - (ry_1 + ry_2)
# If no collision, just skip
if C_x >= 0.0 or C_y >= 0.0:
continue
# Solve collision for the "easier" direction
if C_x > C_y:
grad_C_x_1 = 1.0 if cx_1 - cx_2 >= 0.0 else -1.0
grad_C_x_2 = -1.0 if cx_1 - cx_2 >= 0.0 else 1.0
lagrange = C_x / (grad_C_x_1 * grad_C_x_1 + grad_C_x_2 * grad_C_x_2)
delta_x_1 = -lagrange * grad_C_x_1
delta_x_2 = -lagrange * grad_C_x_2
node_1.location[0] += k * delta_x_1
node_2.location[0] += k * delta_x_2
squared_deltas_sum += k * k * (delta_x_1 * delta_x_1 + delta_x_2 * delta_x_2)
else:
grad_C_y_1 = 1.0 if cy_1 - cy_2 >= 0.0 else -1.0
grad_C_y_2 = -1.0 if cy_1 - cy_2 >= 0.0 else 1.0
lagrange = C_y / (grad_C_y_1 * grad_C_y_1 + grad_C_y_2 * grad_C_y_2)
delta_y_1 = -lagrange * grad_C_y_1
delta_y_2 = -lagrange * grad_C_y_2
node_1.location[1] += k * delta_y_1
node_2.location[1] += k * delta_y_2
squared_deltas_sum += k * k * (delta_y_1 * delta_y_1 + delta_y_2 * delta_y_2)
if verbose:
print("Iteration #" + str(i) + ": " + str(previous_squared_deltas_sum - squared_deltas_sum))
# Check the termination conditiion
if math.fabs(previous_squared_deltas_sum - squared_deltas_sum) < epsilon:
if second_stage:
break
else:
target_space = 0.5 * target_space
second_stage = True
previous_squared_deltas_sum = squared_deltas_sum