nodes/APZmedia-Comfyui-LightDirection-estimation
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Asymetric control of y thresholds

Browse Source
This commit is contained in:
Pablo Apiolazza
2025-10-03 18:53:26 +02:00
parent 6f69851022
commit acafc825cf
3 changed files with 32 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes/nodes.py
+15 -14
View File
@@ -20,7 +20,8 @@ class NormalMapLightEstimator:
"luma_threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
"curve_type": (["linear", "s_curve", "exponential", "logarithmic"], {"default": "s_curve"}),
"x_threshold": ("FLOAT", {"default": 0.4, "min": 0.0, "max": 1.0, "step": 0.01}),
"y_threshold": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0, "step": 0.01}),
"y_threshold_upper": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0, "step": 0.01}),
"y_threshold_lower": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0, "step": 0.01}),
"central_threshold": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01}),
"hard_light_threshold": ("FLOAT", {"default": 0.15, "min": 0.0, "max": 1.0, "step": 0.01}),
"soft_light_threshold": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.01}),
@@ -49,7 +50,7 @@ class NormalMapLightEstimator:
FUNCTION = "estimate_lighting"
def estimate_lighting(self, normal_map, luma_image, luma_threshold, curve_type,
x_threshold, y_threshold, central_threshold,
x_threshold, y_threshold_upper, y_threshold_lower, central_threshold,
hard_light_threshold, soft_light_threshold,
format_mode="auto", normal_standard="OpenGL", analysis_method="advanced", exclusion_mask=None):
"""
@@ -63,7 +64,7 @@ class NormalMapLightEstimator:
luma_processor = LumaMaskProcessor()
light_estimator = CategoricalLightEstimator(
x_threshold, y_threshold, central_threshold,
x_threshold, y_threshold_upper, y_threshold_lower, central_threshold,
hard_light_threshold, soft_light_threshold,
format_mode, normal_standard
)
@@ -117,7 +118,7 @@ class NormalMapLightEstimator:
print(f"Luma image shape: {luma_image.shape}")
print(f"Final mask coverage: {mask.sum().item()}/{mask.numel()} pixels ({mask.float().mean().item():.3f})")
print(f"X threshold: {x_threshold}")
print(f"Y threshold: {y_threshold}")
print(f"Y thresholds: upper={y_threshold_upper}, lower={y_threshold_lower}")
if analysis_method == "legacy":
print("Using LEGACY analysis method (simple mean-based)")
@@ -161,14 +162,14 @@ class NormalMapLightEstimator:
debug_mask = DebugVisualizer.generate_debug_mask(mask)
lit_normals_viz = DebugVisualizer.generate_lit_normals_visualization(normal_map, mask)
# Generate threshold-based classification chart
cluster_delta_chart = DebugVisualizer.create_threshold_classification_chart(normal_map, x_threshold, y_threshold, mask)
cluster_delta_chart = DebugVisualizer.create_threshold_classification_chart(normal_map, x_threshold, y_threshold_upper, y_threshold_lower, mask)
# Generate threshold preview images
x_threshold_preview = self.create_x_threshold_preview(normal_map, x_threshold)
y_threshold_preview = self.create_y_threshold_preview(normal_map, y_threshold)
y_threshold_preview = self.create_y_threshold_preview(normal_map, y_threshold_upper, y_threshold_lower)
print(f"Generated X threshold preview with threshold: {x_threshold}")
print(f"Generated Y threshold preview with threshold: {y_threshold}")
print(f"Generated Y threshold preview with upper: {y_threshold_upper}, lower: {y_threshold_lower}")
return (
x_direction, y_direction, combined_direction,
@@ -287,10 +288,10 @@ class NormalMapLightEstimator:
return tensor_img
@staticmethod
def create_y_threshold_preview(normal_map, y_threshold):
def create_y_threshold_preview(normal_map, y_threshold_upper, y_threshold_lower):
"""
Create Y threshold preview image showing down/center/up zones.
Red = Down, Green = Center, Blue = Up
Create Y threshold preview image showing above/center/below zones.
Red = Above, Green = Center, Blue = Below
"""
import matplotlib.pyplot as plt
import numpy as np
@@ -303,9 +304,9 @@ class NormalMapLightEstimator:
else:
normals_y = normal_map[:, :, :, 1]
# Create classification mask
above_mask = normals_y > y_threshold # Positive Y = surfaces pointing up = light from above
below_mask = normals_y < -y_threshold # Negative Y = surfaces pointing down = light from below
# Create classification mask using separate thresholds
above_mask = normals_y > y_threshold_upper # Positive Y = surfaces pointing up = light from above
below_mask = normals_y < -y_threshold_lower # Negative Y = surfaces pointing down = light from below
center_mask = ~(above_mask | below_mask)
# Create RGB image
@@ -321,7 +322,7 @@ class NormalMapLightEstimator:
plt.switch_backend('Agg')
fig, ax = plt.subplots(figsize=(width/100, height/100))
ax.imshow(preview[0].numpy())
ax.set_title(f'Y Threshold Preview (threshold={y_threshold})', fontsize=12, pad=10)
ax.set_title(f'Y Threshold Preview (upper={y_threshold_upper}, lower={y_threshold_lower})', fontsize=12, pad=10)
# Add legend
legend_elements = [
utils/debug_visualizer.py
+10 -10
View File
@@ -77,7 +77,7 @@ class DebugVisualizer:
return tensor_img
@staticmethod
def create_threshold_classification_chart(normal_map, x_threshold, y_threshold, mask=None):
def create_threshold_classification_chart(normal_map, x_threshold, y_threshold_upper, y_threshold_lower, mask=None):
"""
Create a chart showing threshold-based classification results.
Shows original normal values vs lit-only masked values comparison.
@@ -105,9 +105,9 @@ class DebugVisualizer:
x_right_all = (flat_x > x_threshold).sum().item()
x_total_all = len(flat_x)
y_above_all = (flat_y < -y_threshold).sum().item()
y_center_all = ((flat_y >= -y_threshold) & (flat_y <= y_threshold)).sum().item()
y_below_all = (flat_y > y_threshold).sum().item()
y_above_all = (flat_y > y_threshold_upper).sum().item()
y_center_all = ((flat_y >= -y_threshold_lower) & (flat_y <= y_threshold_upper)).sum().item()
y_below_all = (flat_y < -y_threshold_lower).sum().item()
y_total_all = len(flat_y)
# Calculate lit-only distribution if mask provided
@@ -142,9 +142,9 @@ class DebugVisualizer:
x_right_lit = (lit_x > x_threshold).sum().item()
x_total_lit = len(lit_x)
y_above_lit = (lit_y < -y_threshold).sum().item()
y_center_lit = ((lit_y >= -y_threshold) & (lit_y <= y_threshold)).sum().item()
y_below_lit = (lit_y > y_threshold).sum().item()
y_above_lit = (lit_y > y_threshold_upper).sum().item()
y_center_lit = ((lit_y >= -y_threshold_lower) & (lit_y <= y_threshold_upper)).sum().item()
y_below_lit = (lit_y < -y_threshold_lower).sum().item()
y_total_lit = len(lit_y)
# Create 2x2 subplot layout
@@ -185,7 +185,7 @@ class DebugVisualizer:
y_percentages_all = [count / y_total_all * 100 for count in y_counts_all]
bars3 = ax3.bar(y_categories, y_percentages_all, color=['red', 'green', 'blue'])
ax3.set_title(f'Y Direction - Original Normal Values (threshold={y_threshold})', fontsize=12)
ax3.set_title(f'Y Direction - Original Normal Values (upper={y_threshold_upper}, lower={y_threshold_lower})', fontsize=12)
ax3.set_ylabel('Percentage of All Pixels')
ax3.set_ylim(0, 100)
@@ -199,7 +199,7 @@ class DebugVisualizer:
y_percentages_lit = [count / y_total_lit * 100 for count in y_counts_lit] if y_total_lit > 0 else [0, 0, 0]
bars4 = ax4.bar(y_categories, y_percentages_lit, color=['red', 'green', 'blue'])
ax4.set_title(f'Y Direction - Lit Masked Values (threshold={y_threshold})', fontsize=12)
ax4.set_title(f'Y Direction - Lit Masked Values (upper={y_threshold_upper}, lower={y_threshold_lower})', fontsize=12)
ax4.set_ylabel('Percentage of Lit Pixels')
ax4.set_ylim(0, 100)
@@ -233,7 +233,7 @@ class DebugVisualizer:
y_percentages = [count / y_total_all * 100 for count in y_counts]
bars2 = ax2.bar(y_categories, y_percentages, color=['red', 'green', 'blue'])
ax2.set_title(f'Y Direction - Original Normal Values (threshold={y_threshold})', fontsize=12)
ax2.set_title(f'Y Direction - Original Normal Values (upper={y_threshold_upper}, lower={y_threshold_lower})', fontsize=12)
ax2.set_ylabel('Percentage of Pixels')
ax2.set_ylim(0, 100)
utils/light_estimator.py
+7 -6
View File
@@ -7,12 +7,13 @@ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(
logger = logging.getLogger(__name__)
class CategoricalLightEstimator:
def __init__(self, x_threshold=0.4, y_threshold=0.1, central_threshold=0.3,
def __init__(self, x_threshold=0.4, y_threshold_upper=0.1, y_threshold_lower=0.1, central_threshold=0.3,
hard_light_threshold=0.15, soft_light_threshold=0.35,
format_mode="auto", normal_standard="OpenGL"):
# Initialize thresholds
self.x_threshold = x_threshold
self.y_threshold = y_threshold
self.y_threshold_upper = y_threshold_upper
self.y_threshold_lower = y_threshold_lower
self.central_threshold = central_threshold
self.hard_light_threshold = hard_light_threshold
self.soft_light_threshold = soft_light_threshold
@@ -245,9 +246,9 @@ class CategoricalLightEstimator:
# Y direction classification
y_mean = mean_normal[1].item()
if y_mean > self.y_threshold:
if y_mean > self.y_threshold_upper:
y_category = 'above' # Light from above (surfaces point up)
elif y_mean < -self.y_threshold:
elif y_mean < -self.y_threshold_lower:
y_category = 'below' # Light from below (surfaces point down)
else:
y_category = 'center'
@@ -423,9 +424,9 @@ class CategoricalLightEstimator:
"""
mean_y = y_analysis['mean_y']
if mean_y > self.y_threshold:
if mean_y > self.y_threshold_upper:
return "above" # Light from above (surfaces point up)
elif mean_y < -self.y_threshold:
elif mean_y < -self.y_threshold_lower:
return "below" # Light from below (surfaces point down)
else:
return "center"