Improving light debuggers
This commit is contained in:
Pablo Apiolazza
+29
-8
@@ -25,39 +25,50 @@ class NormalMapLightEstimator:
|
||||
"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}),
|
||||
},
|
||||
"optional": {
|
||||
"format_mode": (["auto", "manual"], {"default": "auto"}),
|
||||
"normal_standard": (["OpenGL", "DirectX", "World_Space", "Object_Space"], {"default": "OpenGL"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = (
|
||||
"STRING", "STRING", "STRING", "FLOAT", "FLOAT", "FLOAT", "FLOAT", "FLOAT",
|
||||
"IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE",
|
||||
"IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE",
|
||||
)
|
||||
|
||||
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", "color_histogram",
|
||||
"debug_mask", "directional_viz", "lit_normals_viz", "colormap_preview",
|
||||
"histogram_before", "histogram_after", "color_histogram",
|
||||
)
|
||||
|
||||
FUNCTION = "estimate_lighting"
|
||||
|
||||
def estimate_lighting(self, normal_map, luma_image, luma_threshold, curve_type,
|
||||
x_threshold, y_threshold, central_threshold,
|
||||
hard_light_threshold, soft_light_threshold):
|
||||
hard_light_threshold, soft_light_threshold,
|
||||
format_mode="auto", normal_standard="OpenGL"):
|
||||
"""
|
||||
Enhanced categorical light direction estimation.
|
||||
"""
|
||||
luma_processor = LumaMaskProcessor()
|
||||
light_estimator = CategoricalLightEstimator(
|
||||
x_threshold, y_threshold, central_threshold,
|
||||
hard_light_threshold, soft_light_threshold
|
||||
hard_light_threshold, soft_light_threshold,
|
||||
format_mode, normal_standard
|
||||
)
|
||||
|
||||
mask = luma_processor.process_with_curves(
|
||||
luma_image, curve_type, None, luma_threshold
|
||||
)
|
||||
|
||||
# Extract lit normals for histogram generation
|
||||
current_mask = mask.bool()
|
||||
lit_normals = normal_map[current_mask]
|
||||
|
||||
results = light_estimator.analyze_directional_categories(normal_map, mask)
|
||||
|
||||
|
||||
x_direction = results['x_category']
|
||||
y_direction = results['y_category']
|
||||
combined_direction = f"{y_direction}-{x_direction}"
|
||||
@@ -66,17 +77,27 @@ class NormalMapLightEstimator:
|
||||
y_confidence = results['confidence']['y_confidence']
|
||||
overall_confidence = results['confidence']['overall_confidence']
|
||||
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()
|
||||
color_histogram = DebugVisualizer.generate_color_histogram(normal_map)
|
||||
|
||||
# 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
|
||||
|
||||
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, color_histogram
|
||||
debug_mask, directional_viz, lit_normals_viz, colormap_preview,
|
||||
histogram_before, histogram_after, color_histogram
|
||||
)
|
||||
|
||||
|
||||
|
||||
+266
-41
@@ -66,58 +66,283 @@ class DebugVisualizer:
|
||||
@staticmethod
|
||||
def generate_colormap_preview(height=256, width=256):
|
||||
"""
|
||||
Generate a visual preview of the colormap used for analysis.
|
||||
Generate a clear 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.
|
||||
"""
|
||||
x = torch.linspace(-1, 1, width)
|
||||
y = torch.linspace(-1, 1, height)
|
||||
xx, yy = torch.meshgrid(y, x, indexing='ij')
|
||||
|
||||
# Mock results for visualization
|
||||
mock_normals = torch.stack([xx, yy, torch.zeros_like(xx)], dim=-1).unsqueeze(0)
|
||||
|
||||
# Generate a color for each point in the grid
|
||||
colormap = torch.zeros(height, width, 3)
|
||||
for i in range(height):
|
||||
for j in range(width):
|
||||
# Pseudo-analysis for color mapping
|
||||
x_val, y_val = xx[i, j], yy[i, j]
|
||||
|
||||
# Determine color based on position
|
||||
red = 128 + 127 * x_val
|
||||
green = 128 + 127 * y_val
|
||||
blue = 128 - 127 * abs(x_val)
|
||||
|
||||
colormap[i, j] = torch.tensor([red, green, blue])
|
||||
|
||||
return colormap.unsqueeze(0).byte().float() / 255.0
|
||||
# Create reference visualization
|
||||
viz = torch.zeros(height, width, 3)
|
||||
|
||||
# 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
|
||||
}
|
||||
|
||||
# Create visual layout
|
||||
section_height = height // 3
|
||||
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),
|
||||
]
|
||||
|
||||
# Draw color sections with labels
|
||||
for label, color, row, col in sections:
|
||||
y_start, x_start = row * section_height, col * section_width
|
||||
|
||||
# Fill section with color
|
||||
viz[y_start:y_start+section_height, x_start:x_start+section_width] = color
|
||||
|
||||
# Add simple label area (could be enhanced with text rendering)
|
||||
label_y = y_start + section_height // 2
|
||||
label_x = x_start + 10
|
||||
|
||||
# Add title and explanation
|
||||
title_area_height = height // 6
|
||||
viz[:title_area_height, :] = 0.8 # Light gray background for title
|
||||
|
||||
return viz.unsqueeze(0)
|
||||
|
||||
@staticmethod
|
||||
def generate_color_histogram(normals, title="Color Distribution"):
|
||||
def generate_color_histogram(normals, title="Normal Distribution"):
|
||||
"""
|
||||
Generate a color histogram visualization of normal vectors.
|
||||
Generate a clear, informative histogram of normal map color distribution.
|
||||
|
||||
Args:
|
||||
normals: Normal vectors tensor (B, H, W, 3)
|
||||
normals: Normal vectors tensor (B, H, W, 3) in [-1, 1] range
|
||||
title: Title for the histogram
|
||||
|
||||
Returns:
|
||||
histogram_viz: RGB visualization of color distribution (1, 256, 256, 3)
|
||||
histogram_viz: Clear visualization with labels and reference (1, 512, 512, 3)
|
||||
"""
|
||||
# Flatten normals to get color distribution
|
||||
flat_normals = normals.view(-1, 3) * 255 + 128 # Scale to 0-255 range
|
||||
flat_normals = torch.clamp(flat_normals, 0, 255)
|
||||
return DebugVisualizer._create_enhanced_histogram(normals, title, show_reference=True)
|
||||
|
||||
# Create histogram visualization
|
||||
@staticmethod
|
||||
def generate_histogram_before(normals, title="Before Masking"):
|
||||
"""
|
||||
Generate histogram of full normal map distribution.
|
||||
|
||||
Args:
|
||||
normals: Normal vectors tensor (B, H, W, 3) for entire image
|
||||
title: Title for the histogram
|
||||
|
||||
Returns:
|
||||
histogram_viz: RGB histogram of full distribution (1, 256, 256, 3)
|
||||
"""
|
||||
return DebugVisualizer._create_color_histogram(normals, title)
|
||||
|
||||
@staticmethod
|
||||
def generate_histogram_after(lit_normals, title="After Masking (Lit Areas)"):
|
||||
"""
|
||||
Generate histogram of lit areas normal distribution.
|
||||
|
||||
Args:
|
||||
lit_normals: Normal vectors tensor (N, 3) for lit areas only
|
||||
title: Title for the histogram
|
||||
|
||||
Returns:
|
||||
histogram_viz: RGB histogram of lit areas distribution (1, 256, 256, 3)
|
||||
"""
|
||||
# Add batch dimension if needed for consistent processing
|
||||
if lit_normals.dim() == 2:
|
||||
lit_normals = lit_normals.unsqueeze(0)
|
||||
|
||||
return DebugVisualizer._create_color_histogram(lit_normals, title)
|
||||
|
||||
@staticmethod
|
||||
def _create_color_histogram(normals, title="Color Distribution"):
|
||||
"""
|
||||
Internal method to create color histogram from normal vectors.
|
||||
|
||||
Args:
|
||||
normals: Normal vectors tensor (B, H, W, 3) or (N, 3)
|
||||
title: Title for the histogram
|
||||
|
||||
Returns:
|
||||
histogram_viz: RGB visualization (1, 256, 256, 3)
|
||||
"""
|
||||
# Flatten normals to analyze distribution
|
||||
flat_normals = normals.view(-1, 3)
|
||||
|
||||
if flat_normals.shape[0] == 0:
|
||||
return torch.zeros(1, 256, 256, 3)
|
||||
|
||||
# Scale normal vectors from [-1, 1] to [0, 255] for histogram bins
|
||||
scaled_normals = (flat_normals + 1.0) * 127.5
|
||||
scaled_normals = torch.clamp(scaled_normals, 0, 255)
|
||||
|
||||
# Create 2D histogram using Red and Green channels as coordinates
|
||||
hist_size = 256
|
||||
histogram = torch.zeros(hist_size, hist_size, 3)
|
||||
histogram = torch.zeros(hist_size, hist_size)
|
||||
|
||||
# Create a simple 2D histogram representation
|
||||
for i in range(min(flat_normals.shape[0], 10000)): # Sample for performance
|
||||
x = int(flat_normals[i, 0].item())
|
||||
y = int(flat_normals[i, 1].item())
|
||||
if 0 <= x < hist_size and 0 <= y < hist_size:
|
||||
# Color based on Z component (blue channel)
|
||||
intensity = flat_normals[i, 2].item() / 255.0
|
||||
histogram[x, y] = torch.tensor([intensity, intensity, 1.0 - intensity])
|
||||
# Sample for performance
|
||||
sample_size = min(flat_normals.shape[0], 100000)
|
||||
indices = torch.randperm(flat_normals.shape[0])[:sample_size]
|
||||
|
||||
return histogram.unsqueeze(0).byte().float() / 255.0
|
||||
for i in indices:
|
||||
normal = scaled_normals[i]
|
||||
r, g, b = normal[0].item(), normal[1].item(), normal[2].item()
|
||||
|
||||
# Use red and green as histogram coordinates
|
||||
r_bin = int(r)
|
||||
g_bin = int(g)
|
||||
|
||||
if 0 <= r_bin < hist_size and 0 <= g_bin < hist_size:
|
||||
histogram[r_bin, g_bin] += 1.0
|
||||
|
||||
# Normalize for visualization
|
||||
if histogram.max() > 0:
|
||||
histogram = histogram / histogram.max()
|
||||
|
||||
# Create RGB visualization
|
||||
rgb_histogram = torch.zeros(hist_size, hist_size, 3)
|
||||
|
||||
for i in range(hist_size):
|
||||
for j in range(hist_size):
|
||||
density = histogram[i, j].item()
|
||||
|
||||
if density > 0:
|
||||
# Color by position and density
|
||||
# X direction (red) + Y direction (green) + density (blue)
|
||||
x_intensity = i / 255.0 # Red channel shows X position
|
||||
y_intensity = j / 255.0 # Green channel shows Y position
|
||||
z_intensity = density # Blue channel shows density
|
||||
|
||||
rgb_histogram[i, j] = torch.tensor([
|
||||
x_intensity, # X direction indicator
|
||||
y_intensity, # Y direction indicator
|
||||
z_intensity # Density indicator
|
||||
])
|
||||
|
||||
return rgb_histogram.unsqueeze(0)
|
||||
|
||||
@staticmethod
|
||||
def _create_enhanced_histogram(normals, title="Normal Distribution", show_reference=True):
|
||||
"""
|
||||
Create an enhanced histogram with labels, legends, and clear visual design.
|
||||
|
||||
Args:
|
||||
normals: Normal vectors tensor (B, H, W, 3) or (N, 3)
|
||||
title: Title for the histogram
|
||||
show_reference: Whether to show reference information
|
||||
|
||||
Returns:
|
||||
enhanced_viz: Complete visualization with labels (1, 512, 512, 3)
|
||||
"""
|
||||
# Get base histogram (256x256)
|
||||
base_histogram = DebugVisualizer._create_color_histogram_internal(normals)
|
||||
|
||||
if base_histogram is None:
|
||||
return torch.zeros(1, 512, 512, 3)
|
||||
|
||||
# Create enhanced visualization with labels and reference
|
||||
viz_size = 512
|
||||
hist_size = 256
|
||||
enhanced = torch.zeros(viz_size, viz_size, 3)
|
||||
|
||||
# Place histogram in top-left area (256x256)
|
||||
enhanced[:hist_size, :hist_size] = base_histogram[0]
|
||||
|
||||
# Add reference information (right side)
|
||||
if show_reference:
|
||||
enhanced = DebugVisualizer._add_histogram_labels(enhanced, title)
|
||||
|
||||
return enhanced.unsqueeze(0)
|
||||
|
||||
@staticmethod
|
||||
def _create_color_histogram_internal(normals):
|
||||
"""Internal histogram creation without labels"""
|
||||
flat_normals = normals.view(-1, 3)
|
||||
if flat_normals.shape[0] == 0:
|
||||
return None
|
||||
|
||||
# Scale to [0, 255] for histogram bins
|
||||
scaled_normals = (flat_normals + 1.0) * 127.5
|
||||
scaled_normals = torch.clamp(scaled_normals, 0, 255)
|
||||
|
||||
# Create histogram
|
||||
hist_size = 256
|
||||
histogram = torch.zeros(hist_size, hist_size)
|
||||
|
||||
sample_size = min(flat_normals.shape[0], 50000)
|
||||
indices = torch.randperm(flat_normals.shape[0])[:sample_size]
|
||||
|
||||
for i in indices:
|
||||
normal = scaled_normals[i]
|
||||
r, g, b = normal[0].item(), normal[1].item(), normal[2].item()
|
||||
r_bin, g_bin = int(r), int(g)
|
||||
|
||||
if 0 <= r_bin < hist_size and 0 <= g_bin < hist_size:
|
||||
histogram[r_bin, g_bin] += 1.0
|
||||
|
||||
# Normalize
|
||||
if histogram.max() > 0:
|
||||
histogram = histogram / histogram.max()
|
||||
|
||||
# Create visualization with clear color mapping
|
||||
rgb_histogram = torch.zeros(hist_size, hist_size, 3)
|
||||
|
||||
# Add coordinate grid lines
|
||||
grid_color = torch.tensor([0.3, 0.3, 0.3])
|
||||
for i in range(0, hist_size, 32):
|
||||
rgb_histogram[i, :, :] = grid_color
|
||||
rgb_histogram[:, i, :] = grid_color
|
||||
|
||||
# Add density visualization
|
||||
for i in range(hist_size):
|
||||
for j in range(hist_size):
|
||||
density = histogram[i, j].item()
|
||||
if density > 0.01: # Threshold to reduce noise
|
||||
# Heatmap: blue (low) to red (high)
|
||||
if density < 0.33:
|
||||
intensity = density * 3
|
||||
rgb_histogram[i, j] = torch.tensor([0.0, 0.0, intensity])
|
||||
elif density < 0.66:
|
||||
intensity = (density - 0.33) * 3
|
||||
rgb_histogram[i, j] = torch.tensor([intensity, 0.0, 0.0])
|
||||
else:
|
||||
intensity = (density - 0.66) * 3
|
||||
rgb_histogram[i, j] = torch.tensor([1.0, intensity, 0.0])
|
||||
|
||||
return rgb_histogram.unsqueeze(0)
|
||||
|
||||
@staticmethod
|
||||
def _add_histogram_labels(enhanced_viz, title):
|
||||
"""Add labels and reference information to histogram"""
|
||||
viz_size = enhanced_viz.shape[0]
|
||||
|
||||
# Add title (top area)
|
||||
title_y = 20
|
||||
# Simple title for now - could be enhanced with text rendering
|
||||
|
||||
# Add axis labels
|
||||
# X-axis: "Left ← X Direction → Right"
|
||||
# Y-axis: "Bottom ← Y Direction → Top"
|
||||
|
||||
# Add color legend (right side)
|
||||
legend_x = int(viz_size * 0.8)
|
||||
legend_width = viz_size - legend_x
|
||||
|
||||
# Color legend: show what colors mean
|
||||
legend_colors = [
|
||||
torch.tensor([0.0, 0.0, 1.0]), # Blue = low density
|
||||
torch.tensor([0.0, 1.0, 0.0]), # Green = medium density
|
||||
torch.tensor([1.0, 0.0, 0.0]), # Red = high density
|
||||
]
|
||||
|
||||
for i, color in enumerate(legend_colors):
|
||||
y = int(viz_size * 0.3) + i * 30
|
||||
enhanced_viz[y:y+20, legend_x:legend_x+20] = color
|
||||
|
||||
return enhanced_viz
|
||||
|
||||
@@ -6,15 +6,21 @@ class CategoricalLightEstimator:
|
||||
"""
|
||||
|
||||
def __init__(self, x_threshold=0.1, y_threshold=0.1, central_threshold=0.3,
|
||||
hard_light_threshold=0.15, soft_light_threshold=0.35):
|
||||
hard_light_threshold=0.15, soft_light_threshold=0.35,
|
||||
format_mode="auto", normal_standard="OpenGL"):
|
||||
# Directional thresholds
|
||||
self.x_threshold = x_threshold
|
||||
self.y_threshold = y_threshold
|
||||
self.central_threshold = central_threshold
|
||||
|
||||
|
||||
# Hard/Soft classification thresholds
|
||||
self.hard_light_threshold = hard_light_threshold
|
||||
self.soft_light_threshold = soft_light_threshold
|
||||
|
||||
# Normal map format settings
|
||||
self.format_mode = format_mode
|
||||
self.normal_standard = normal_standard
|
||||
self.detected_format = None
|
||||
|
||||
def analyze_directional_categories(self, normals, mask):
|
||||
"""
|
||||
@@ -34,6 +40,10 @@ class CategoricalLightEstimator:
|
||||
mask = mask.squeeze(-1)
|
||||
current_mask = mask.bool()
|
||||
|
||||
# Detect and validate normal map format
|
||||
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)
|
||||
|
||||
@@ -49,6 +59,18 @@ class CategoricalLightEstimator:
|
||||
# Analyze color distribution after masking
|
||||
color_analysis_after = self._analyze_color_distribution(lit_normals)
|
||||
|
||||
# DEBUG: Print key information for pipeline verification
|
||||
print("=== Pipeline Verification ===")
|
||||
print(f"Input normals shape: {normals.shape}")
|
||||
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']}")
|
||||
print("==========================")
|
||||
|
||||
# Extract XY components
|
||||
xy_normals = lit_normals[:, :2]
|
||||
|
||||
@@ -247,3 +269,48 @@ class CategoricalLightEstimator:
|
||||
'confidence': {'x_confidence': 0.0, 'y_confidence': 0.0, 'quality_confidence': 0.0, 'overall_confidence': 0.0},
|
||||
'quality_analysis': {'spread': 0.0, 'confidence': 0.0}
|
||||
}
|
||||
|
||||
def _detect_normal_format(self, normals):
|
||||
"""
|
||||
Auto-detect the normal map format by analyzing color distribution patterns.
|
||||
|
||||
Args:
|
||||
normals: Normal vectors tensor (B, H, W, 3)
|
||||
|
||||
Returns:
|
||||
detected_format: String indicating the likely format
|
||||
"""
|
||||
# Flatten to analyze color distribution
|
||||
flat_normals = normals.view(-1, 3)
|
||||
|
||||
# Calculate mean colors
|
||||
mean_colors = torch.mean(flat_normals, dim=0)
|
||||
|
||||
# Check for DirectX vs OpenGL pattern
|
||||
# DirectX typically has more variation in green channel if Y is inverted
|
||||
green_std = torch.std(flat_normals[:, 1]).item()
|
||||
blue_std = torch.std(flat_normals[:, 2]).item()
|
||||
|
||||
# Simple heuristic: if green channel has significantly more variation than blue,
|
||||
# it might be DirectX format (where Y is inverted)
|
||||
if green_std > blue_std * 1.5:
|
||||
return "DirectX"
|
||||
else:
|
||||
return "OpenGL"
|
||||
|
||||
def _interpret_normal_colors(self, x_val, y_val, z_val):
|
||||
"""
|
||||
Interpret normal vector values based on the detected format.
|
||||
|
||||
Args:
|
||||
x_val, y_val, z_val: Color components from normal map
|
||||
|
||||
Returns:
|
||||
interpreted_x, interpreted_y: Properly oriented normal components
|
||||
"""
|
||||
if self.detected_format == "DirectX":
|
||||
# DirectX inverts Y axis
|
||||
return x_val, -y_val
|
||||
else:
|
||||
# OpenGL, World Space, Object Space (standard orientation)
|
||||
return x_val, y_val
|
||||
|
||||
Reference in New Issue
Block a user