From 06f8a420ef70399bc866f6381688c59d5e01bf2f Mon Sep 17 00:00:00 2001 From: Pablo Apiolazza <31888952+APZmedia@users.noreply.github.com> Date: Fri, 3 Oct 2025 10:16:06 +0200 Subject: [PATCH] Improving light debuggers --- nodes/nodes.py | 37 ++++- utils/debug_visualizer.py | 307 +++++++++++++++++++++++++++++++++----- utils/light_estimator.py | 71 ++++++++- 3 files changed, 364 insertions(+), 51 deletions(-) diff --git a/nodes/nodes.py b/nodes/nodes.py index 8bdb2f3..671fa22 100644 --- a/nodes/nodes.py +++ b/nodes/nodes.py @@ -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 ) diff --git a/utils/debug_visualizer.py b/utils/debug_visualizer.py index b1c3d49..fbd6cbf 100644 --- a/utils/debug_visualizer.py +++ b/utils/debug_visualizer.py @@ -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 diff --git a/utils/light_estimator.py b/utils/light_estimator.py index 4e2cf41..99a307b 100644 --- a/utils/light_estimator.py +++ b/utils/light_estimator.py @@ -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