diff --git a/nodes/nodes.py b/nodes/nodes.py index 671fa22..d5b4970 100644 --- a/nodes/nodes.py +++ b/nodes/nodes.py @@ -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 ) diff --git a/test-images/Front-left.png b/test-images/Front-left.png new file mode 100644 index 0000000..bf6147d Binary files /dev/null and b/test-images/Front-left.png differ diff --git a/test-images/Soft-right-2.png b/test-images/Soft-right-2.png new file mode 100644 index 0000000..1943749 Binary files /dev/null and b/test-images/Soft-right-2.png differ diff --git a/test-images/Soft-right.png b/test-images/Soft-right.png new file mode 100644 index 0000000..cf05529 Binary files /dev/null and b/test-images/Soft-right.png differ diff --git a/utils/debug_visualizer.py b/utils/debug_visualizer.py index fbd6cbf..0425212 100644 --- a/utils/debug_visualizer.py +++ b/utils/debug_visualizer.py @@ -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""" diff --git a/utils/light_estimator.py b/utils/light_estimator.py index 99a307b..9c37a0f 100644 --- a/utils/light_estimator.py +++ b/utils/light_estimator.py @@ -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.