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

Clustering, examples an debug visualization

Browse Source
This commit is contained in:
Pablo Apiolazza
2025-10-03 11:44:16 +02:00
parent 06f8a420ef
commit f036da13ac
6 changed files with 339 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes/nodes.py
+8 -14
View File
@@ -39,8 +39,8 @@ class NormalMapLightEstimator:
RETURN_NAMES = (
"x_direction", "y_direction", "combined_direction", "hard_soft_index",
"x_confidence", "y_confidence", "overall_confidence", "spread_value",
"debug_mask", "directional_viz", "lit_normals_viz", "colormap_preview",
"histogram_before", "histogram_after", "color_histogram",
"debug_mask", "lit_normals_viz", "colormap_preview",
"directional_cluster_viz", "cluster_distribution_chart", "lighting_summary_viz",
)
FUNCTION = "estimate_lighting"
@@ -79,25 +79,19 @@ class NormalMapLightEstimator:
spread_value = results['quality_analysis']['spread']
debug_mask = DebugVisualizer.generate_debug_mask(mask)
directional_viz = DebugVisualizer.generate_directional_visualization(normal_map, results)
lit_normals_viz = DebugVisualizer.generate_lit_normals_visualization(normal_map, mask)
colormap_preview = DebugVisualizer.generate_colormap_preview()
# Generate separate histograms for before/after masking
histogram_before = DebugVisualizer.generate_histogram_before(normal_map)
# Ensure lit_normals is in correct format for histogram generation
if lit_normals.dim() == 2:
lit_normals_4d = lit_normals.unsqueeze(0) # Add batch dimension
else:
lit_normals_4d = lit_normals
histogram_after = DebugVisualizer.generate_histogram_after(lit_normals_4d)
color_histogram = DebugVisualizer.generate_color_histogram(normal_map) # Combined view
# Generate new debug visualizations
directional_cluster_viz = DebugVisualizer.generate_directional_cluster_viz(normal_map, results)
cluster_distribution_chart = DebugVisualizer.generate_cluster_distribution_chart(results)
lighting_summary_viz = DebugVisualizer.generate_lighting_summary_viz(results)
return (
x_direction, y_direction, combined_direction, hard_soft_index,
x_confidence, y_confidence, overall_confidence, spread_value,
debug_mask, directional_viz, lit_normals_viz, colormap_preview,
histogram_before, histogram_after, color_histogram
debug_mask, lit_normals_viz, colormap_preview,
directional_cluster_viz, cluster_distribution_chart, lighting_summary_viz
)
test-images/Front-left.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 2.1 MiB

test-images/Soft-right-2.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 2.1 MiB

test-images/Soft-right.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 21 MiB

utils/debug_visualizer.py
+129 -69
View File
@@ -28,89 +28,99 @@ class DebugVisualizer:
lit_normals_viz = normal_rgb * lit_mask
return lit_normals_viz
@staticmethod
def generate_directional_visualization(normals, results):
"""
Generate color-coded visualization of directional analysis.
"""
batch_size, height, width, _ = normals.shape
viz = torch.zeros(batch_size, height, width, 3, device=normals.device)
x_category = results['x_category']
y_category = results['y_category']
hard_soft_index = results['hard_soft_index']
overall_confidence = results['confidence']['overall_confidence']
# Base colors
if x_category == "left":
base_color = torch.tensor([255, 100, 100], device=normals.device)
elif x_category == "right":
base_color = torch.tensor([200, 80, 80], device=normals.device)
else:
base_color = torch.tensor([150, 150, 150], device=normals.device)
# Adjust for Y direction
if y_category == "top":
base_color += torch.tensor([0, 100, 50], device=normals.device)
elif y_category == "bottom":
base_color += torch.tensor([50, -50, 0], device=normals.device)
# Apply softness and confidence
softness_factor = 1.0 - hard_soft_index
final_color = base_color * (0.7 + 0.3 * softness_factor) * overall_confidence
final_color = torch.clamp(final_color, 0, 255)
viz[:, :, :] = final_color.byte()
return viz.float() / 255.0
@staticmethod
def generate_colormap_preview(height=256, width=256):
def generate_colormap_preview(height=512, width=512):
"""
Generate a clear reference guide for normal map color encoding.
Generate a comprehensive reference guide for normal map color encoding.
This shows how 3D surface orientations are encoded as RGB colors
in normal maps, making it easy to understand what each color means.
This creates an educational reference showing exactly how 3D surface
orientations are encoded as RGB colors in normal maps.
"""
# Create reference visualization
viz = torch.zeros(height, width, 3)
# Create main visualization canvas
viz = torch.ones(height, width, 3) * 0.9 # Light gray background
# Color legend for normal map encoding
normal_colors = {
'right': torch.tensor([1.0, 0.5, 0.5]), # Red - Right facing
'left': torch.tensor([0.5, 1.0, 1.0]), # Cyan - Left facing
'up': torch.tensor([0.5, 1.0, 0.5]), # Green - Up facing
'down': torch.tensor([1.0, 0.5, 1.0]), # Purple - Down facing
'out': torch.tensor([0.5, 0.5, 1.0]), # Blue - Outward facing
'flat': torch.tensor([0.5, 0.5, 0.5]), # Gray - Flat/neutral
}
# Title section
title_height = 60
viz[:title_height, :] = torch.tensor([0.7, 0.7, 0.7]) # Darker background for title
# Create visual layout
section_height = height // 3
# Color reference sections in a 2x3 grid
section_height = (height - title_height - 40) // 2 # Leave space for labels
section_width = width // 3
sections = [
('X+ (Right)', normal_colors['right'], 0, 0),
('X- (Left)', normal_colors['left'], 0, 1),
('Y+ (Up)', normal_colors['up'], 1, 0),
('Y- (Down)', normal_colors['down'], 1, 1),
('Z+ (Out)', normal_colors['out'], 2, 0),
('Z=0 (Flat)', normal_colors['flat'], 2, 1),
# Define normal map color encodings with clear labels
normal_directions = [
# Row 1: X-axis directions
{
'label': 'X+ (Right)',
'description': 'Surfaces facing RIGHT are RED',
'color': torch.tensor([1.0, 0.5, 0.5]), # Red
'position': 'Lit by RIGHT-side lighting',
'row': 0, 'col': 0
},
{
'label': 'X- (Left)',
'description': 'Surfaces facing LEFT are CYAN',
'color': torch.tensor([0.5, 1.0, 1.0]), # Cyan
'position': 'Lit by LEFT-side lighting',
'row': 0, 'col': 1
},
{
'label': 'Y+ (Up)',
'description': 'Surfaces facing UP are GREEN',
'color': torch.tensor([0.5, 1.0, 0.5]), # Green
'position': 'Lit by TOP lighting',
'row': 0, 'col': 2
},
# Row 2: Y-axis and Z-axis directions
{
'label': 'Y- (Down)',
'description': 'Surfaces facing DOWN are PURPLE',
'color': torch.tensor([1.0, 0.5, 1.0]), # Purple
'position': 'Lit by BOTTOM lighting',
'row': 1, 'col': 0
},
{
'label': 'Z+ (Out)',
'description': 'Surfaces facing OUT are BLUE',
'color': torch.tensor([0.5, 0.5, 1.0]), # Blue
'position': 'Lit by FRONT lighting',
'row': 1, 'col': 1
},
{
'label': 'Z=0 (Flat)',
'description': 'Flat surfaces are GRAY',
'color': torch.tensor([0.5, 0.5, 0.5]), # Gray
'position': 'Rarely lit directly',
'row': 1, 'col': 2
}
]
# Draw color sections with labels
for label, color, row, col in sections:
y_start, x_start = row * section_height, col * section_width
# Draw each section
for direction in normal_directions:
row, col = direction['row'], direction['col']
y_start = title_height + row * section_height
x_start = col * section_width
# Fill section with color
viz[y_start:y_start+section_height, x_start:x_start+section_width] = color
# Fill section with the characteristic color
color_block_size = 80
y_center = y_start + section_height // 2 - color_block_size // 2
x_center = x_start + section_width // 2 - color_block_size // 2
# Add simple label area (could be enhanced with text rendering)
label_y = y_start + section_height // 2
label_x = x_start + 10
viz[y_center:y_center+color_block_size, x_center:x_center+color_block_size] = direction['color']
# Add title and explanation
title_area_height = height // 6
viz[:title_area_height, :] = 0.8 # Light gray background for title
# Add section border
border_color = torch.tensor([0.3, 0.3, 0.3])
viz[y_start:y_start+section_height, x_start:x_start+5] = border_color
viz[y_start:y_start+section_height, x_start+section_width-5:x_start+section_width] = border_color
viz[y_start:y_start+5, x_start:x_start+section_width] = border_color
viz[y_start+section_height-5:y_start+section_height, x_start:x_start+section_width] = border_color
# Add usage examples at bottom
example_height = 100
example_y = height - example_height
viz[example_y:, :] = torch.tensor([0.95, 0.95, 0.95])
return viz.unsqueeze(0)
@@ -260,6 +270,56 @@ class DebugVisualizer:
return enhanced.unsqueeze(0)
@staticmethod
def generate_directional_cluster_viz(normals, results):
"""
Visualize directional clusters on the normal map.
"""
if 'clustering_before' not in results:
return torch.zeros_like(normals)
# Create a blank canvas
viz = torch.zeros_like(normals)
# Define colors for each cluster
cluster_colors = {
'right': torch.tensor([1.0, 0.0, 0.0]), # Red
'left': torch.tensor([0.0, 1.0, 1.0]), # Cyan
'up': torch.tensor([0.0, 1.0, 0.0]), # Green
'down': torch.tensor([1.0, 0.0, 1.0]), # Purple
'front': torch.tensor([0.0, 0.0, 1.0]), # Blue
'flat': torch.tensor([0.5, 0.5, 0.5]), # Gray
}
# Iterate through pixels and color them based on their cluster
# This is computationally expensive, so it should be used for debugging
# A more optimized version could use masks for each cluster
return viz.unsqueeze(0)
@staticmethod
def generate_cluster_distribution_chart(results):
"""
Create a bar chart of lit vs unlit distribution for each cluster.
"""
# (Implementation to be added)
return torch.zeros(1, 512, 512, 3)
@staticmethod
def generate_lighting_summary_viz(results):
"""
Create a summary infographic of the final results.
"""
# (Implementation to be added)
return torch.zeros(1, 256, 512, 3)
@staticmethod
def generate_lighting_summary_viz(results):
"""
Create a summary infographic of the final results.
"""
# (Implementation to be added)
return torch.zeros(1, 256, 512, 3)
@staticmethod
def _create_color_histogram_internal(normals):
"""Internal histogram creation without labels"""
utils/light_estimator.py
+202 -11
View File
@@ -44,20 +44,25 @@ class CategoricalLightEstimator:
detected_format = self._detect_normal_format(normals) if self.format_mode == "auto" else self.normal_standard
self.detected_format = detected_format
# Analyze color distribution before masking
color_analysis_before = self._analyze_color_distribution(normals)
# Analyze directional clustering before masking
clustering_before = self._analyze_directional_clustering(normals, None)
# Apply mask to normals
lit_normals = normals[current_mask]
if lit_normals.numel() == 0:
return self._empty_categories_with_color_analysis(batch_size, normals.device, color_analysis_before)
return self._empty_categories_with_clustering(batch_size, normals.device, clustering_before)
if lit_normals.dim() == 1:
lit_normals = lit_normals.unsqueeze(0)
# Analyze color distribution after masking
color_analysis_after = self._analyze_color_distribution(lit_normals)
# Analyze directional clustering after masking
clustering_after = self._analyze_directional_clustering(lit_normals.unsqueeze(0), current_mask.unsqueeze(0))
# Enhanced analysis using directional clustering
enhanced_results = self._analyze_directional_clustering_enhanced(
normals, current_mask, clustering_before, clustering_after
)
# DEBUG: Print key information for pipeline verification
print("=== Pipeline Verification ===")
@@ -65,10 +70,18 @@ class CategoricalLightEstimator:
print(f"Mask shape: {mask.shape}")
print(f"Lit normals count: {lit_normals.shape[0]}")
print(f"Total pixels: {normals.shape[1] * normals.shape[2]}")
print(f"Color before - Mean: {color_analysis_before['mean_color']}")
print(f"Color before - Dominant: {color_analysis_before['dominant_direction']}")
print(f"Color after - Mean: {color_analysis_after['mean_color']}")
print(f"Color after - Dominant: {color_analysis_after['dominant_direction']}")
# Show clustering analysis
print("=== Directional Clustering ===")
for cluster_name, cluster_data in clustering_before['clusters'].items():
before_pct = cluster_data['percentage']
after_pct = clustering_after['clusters'][cluster_name]['percentage']
effect = after_pct - before_pct
print(f"{cluster_name}: Before={before_pct:.3f}, After={after_pct:.3f}, Effect={effect:.3f}")
print(f"Dominant cluster before: {clustering_before['dominant_cluster']}")
print(f"Dominant cluster after: {clustering_after['dominant_cluster']}")
print(f"Enhanced analysis: {enhanced_results['primary_direction']}")
print("==========================")
# Extract XY components
@@ -96,14 +109,20 @@ class CategoricalLightEstimator:
'overall_confidence': (x_analysis['confidence'] + y_analysis['confidence'] + quality_analysis['confidence']) / 3
}
# Enhanced analysis using directional clustering
enhanced_results = self._analyze_directional_clustering_enhanced(
normals, current_mask, clustering_before, clustering_after
)
return {
'x_category': x_category,
'y_category': y_category,
'hard_soft_index': hard_soft_index,
'confidence': confidence,
'quality_analysis': quality_analysis,
'color_analysis_before': color_analysis_before,
'color_analysis_after': color_analysis_after,
'clustering_before': clustering_before,
'clustering_after': clustering_after,
'enhanced_analysis': enhanced_results,
'lit_pixel_count': lit_normals.shape[0],
'total_pixel_count': normals.shape[1] * normals.shape[2]
}
@@ -270,6 +289,178 @@ class CategoricalLightEstimator:
'quality_analysis': {'spread': 0.0, 'confidence': 0.0}
}
def _analyze_directional_clustering(self, normals, mask=None):
"""
Analyze normal vectors by clustering them into directional categories.
Args:
normals: Normal vectors tensor (B, H, W, 3)
mask: Optional binary mask to limit analysis
Returns:
clustering_results: Dictionary with directional cluster analysis
"""
if normals.dim() == 4:
flat_normals = normals.view(-1, 3)
height, width = normals.shape[1], normals.shape[2]
else:
flat_normals = normals
height, width = 1, normals.shape[0]
if mask is not None:
if mask.dim() == 4:
mask = mask.squeeze(-1)
mask_flat = mask.view(-1)
flat_normals = flat_normals[mask_flat.bool()]
if flat_normals.shape[0] == 0:
return self._empty_directional_clustering()
# Define directional color clusters
clusters = self._define_directional_clusters()
# Analyze each cluster
cluster_results = {}
for cluster_name, color_ranges in clusters.items():
cluster_normals = self._find_normals_in_clusters(flat_normals, color_ranges)
cluster_results[cluster_name] = {
'count': cluster_normals.shape[0],
'percentage': cluster_normals.shape[0] / flat_normals.shape[0],
'mean_normal': torch.mean(cluster_normals, dim=0) if cluster_normals.shape[0] > 0 else torch.zeros(3)
}
return {
'clusters': cluster_results,
'total_normals': flat_normals.shape[0],
'dominant_cluster': max(cluster_results.keys(), key=lambda k: cluster_results[k]['count'])
}
def _define_directional_clusters(self):
"""
Define color ranges for each directional cluster.
Returns:
clusters: Dictionary mapping direction names to color ranges
"""
# Define color ranges for each direction (in RGB space)
clusters = {
'right': { # Red-dominant colors (X+)
'x_range': [0.4, 1.0], # Red channel high
'y_range': [0.3, 0.7], # Green channel mid
'z_range': [0.3, 0.7] # Blue channel mid
},
'left': { # Cyan-dominant colors (X-)
'x_range': [0.0, 0.6], # Red channel low
'y_range': [0.4, 1.0], # Green channel high
'z_range': [0.4, 1.0] # Blue channel high
},
'up': { # Green-dominant colors (Y+)
'x_range': [0.3, 0.7], # Red channel mid
'y_range': [0.4, 1.0], # Green channel high
'z_range': [0.3, 0.7] # Blue channel mid
},
'down': { # Purple-dominant colors (Y-)
'x_range': [0.4, 1.0], # Red channel high
'y_range': [0.0, 0.6], # Green channel low
'z_range': [0.4, 1.0] # Blue channel high
},
'front': { # Blue-dominant colors (Z+)
'x_range': [0.3, 0.7], # Red channel mid
'y_range': [0.3, 0.7], # Green channel mid
'z_range': [0.4, 1.0] # Blue channel high
},
'flat': { # Gray/neutral colors (Z≈0)
'x_range': [0.3, 0.7], # Red channel mid
'y_range': [0.3, 0.7], # Green channel mid
'z_range': [0.0, 0.6] # Blue channel low
}
}
return clusters
def _find_normals_in_clusters(self, normals, color_ranges):
"""
Find normal vectors that fall within specified color ranges.
Args:
normals: Normal vectors tensor (N, 3)
color_ranges: Dictionary with x_range, y_range, z_range
Returns:
filtered_normals: Normals within the color ranges
"""
x_min, x_max = color_ranges['x_range']
y_min, y_max = color_ranges['y_range']
z_min, z_max = color_ranges['z_range']
# Filter normals based on color ranges
mask = (normals[:, 0] >= x_min) & (normals[:, 0] <= x_max) & \
(normals[:, 1] >= y_min) & (normals[:, 1] <= y_max) & \
(normals[:, 2] >= z_min) & (normals[:, 2] <= z_max)
return normals[mask]
def _analyze_directional_clustering_enhanced(self, normals, mask, clustering_before, clustering_after):
"""
Enhanced analysis using directional clustering with luminance correlation.
Args:
normals: Full normal map tensor
mask: Binary mask tensor
clustering_before: Clustering results for full image
clustering_after: Clustering results for lit areas
Returns:
enhanced_results: Advanced lighting analysis
"""
# Calculate lighting effect on each directional cluster
lighting_effects = {}
for cluster_name in clustering_before['clusters'].keys():
before_pct = clustering_before['clusters'][cluster_name]['percentage']
after_pct = clustering_after['clusters'][cluster_name]['percentage']
# Calculate how much this cluster is affected by lighting
lighting_effect = after_pct - before_pct
lighting_effects[cluster_name] = {
'before_percentage': before_pct,
'after_percentage': after_pct,
'lighting_effect': lighting_effect,
'is_illuminated': lighting_effect > 0.05 # Threshold for significant illumination
}
# Determine primary lighting direction based on cluster effects
illuminated_clusters = [name for name, effect in lighting_effects.items() if effect['is_illuminated']]
# Map cluster illumination to lighting direction
direction_mapping = {
'right': 'right',
'left': 'left',
'up': 'top',
'down': 'bottom',
'front': 'front'
}
detected_directions = [direction_mapping.get(cluster, 'unknown') for cluster in illuminated_clusters]
return {
'lighting_effects': lighting_effects,
'illuminated_clusters': illuminated_clusters,
'detected_directions': detected_directions,
'primary_direction': max(detected_directions, default='unknown') if detected_directions else 'diffuse'
}
def _empty_directional_clustering(self):
"""Return empty clustering results"""
clusters = {name: {'count': 0, 'percentage': 0.0, 'mean_normal': torch.zeros(3)}
for name in ['right', 'left', 'up', 'down', 'front', 'flat']}
return {
'clusters': clusters,
'total_normals': 0,
'dominant_cluster': 'none'
}
def _detect_normal_format(self, normals):
"""
Auto-detect the normal map format by analyzing color distribution patterns.