diff --git a/nodes/nodes.py b/nodes/nodes.py index b18fbef..490c6b8 100644 --- a/nodes/nodes.py +++ b/nodes/nodes.py @@ -125,8 +125,8 @@ class NormalMapLightEstimator: debug_mask = DebugVisualizer.generate_debug_mask(mask) lit_normals_viz = DebugVisualizer.generate_lit_normals_visualization(normal_map, mask) - # Generate cluster delta visualization - cluster_delta_chart = DebugVisualizer.create_cluster_delta_chart(results) + # Generate threshold-based classification chart + cluster_delta_chart = DebugVisualizer.create_threshold_classification_chart(normal_map, x_threshold, y_threshold) # Generate threshold preview images x_threshold_preview = self.create_x_threshold_preview(normal_map, x_threshold) diff --git a/utils/debug_visualizer.py b/utils/debug_visualizer.py index 8cf3a1f..bfadd5c 100644 --- a/utils/debug_visualizer.py +++ b/utils/debug_visualizer.py @@ -75,3 +75,82 @@ class DebugVisualizer: # Ensure proper format for ComfyUI tensor_img = tensor_img.clamp(0.0, 1.0).type(torch.float32) return tensor_img + + @staticmethod + def create_threshold_classification_chart(normal_map, x_threshold, y_threshold): + """ + Create a chart showing threshold-based classification results. + """ + import torch + import numpy as np + + plt.switch_backend('Agg') + fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6)) + + # Convert normal map to -1 to +1 range if needed + if normal_map.max() <= 1.0: + normals_x = (normal_map[:, :, :, 0] * 2.0) - 1.0 + normals_y = (normal_map[:, :, :, 1] * 2.0) - 1.0 + else: + normals_x = normal_map[:, :, :, 0] + normals_y = normal_map[:, :, :, 1] + + # Flatten for analysis + flat_x = normals_x.view(-1) + flat_y = normals_y.view(-1) + + # X Direction Classification + x_left = (flat_x < -x_threshold).sum().item() + x_center = ((flat_x >= -x_threshold) & (flat_x <= x_threshold)).sum().item() + x_right = (flat_x > x_threshold).sum().item() + x_total = len(flat_x) + + x_categories = ['Light from Right', 'Center', 'Light from Left'] + x_counts = [x_left, x_center, x_right] + x_percentages = [count / x_total * 100 for count in x_counts] + + bars1 = ax1.bar(x_categories, x_percentages, color=['red', 'green', 'blue']) + ax1.set_title(f'X Direction Classification (threshold={x_threshold})', fontsize=12) + ax1.set_ylabel('Percentage of Pixels') + ax1.set_ylim(0, 100) + + # Add percentage labels on bars + for bar, pct in zip(bars1, x_percentages): + height = bar.get_height() + ax1.text(bar.get_x() + bar.get_width()/2., height + 1, + f'{pct:.1f}%', ha='center', va='bottom') + + # Y Direction Classification + y_above = (flat_y < -y_threshold).sum().item() + y_center = ((flat_y >= -y_threshold) & (flat_y <= y_threshold)).sum().item() + y_below = (flat_y > y_threshold).sum().item() + y_total = len(flat_y) + + y_categories = ['Light from Above', 'Center', 'Light from Below'] + y_counts = [y_above, y_center, y_below] + y_percentages = [count / y_total * 100 for count in y_counts] + + bars2 = ax2.bar(y_categories, y_percentages, color=['red', 'green', 'blue']) + ax2.set_title(f'Y Direction Classification (threshold={y_threshold})', fontsize=12) + ax2.set_ylabel('Percentage of Pixels') + ax2.set_ylim(0, 100) + + # Add percentage labels on bars + for bar, pct in zip(bars2, y_percentages): + height = bar.get_height() + ax2.text(bar.get_x() + bar.get_width()/2., height + 1, + f'{pct:.1f}%', ha='center', va='bottom') + + plt.tight_layout() + + # Convert to tensor + buf = io.BytesIO() + plt.savefig(buf, format='png', dpi=100, bbox_inches='tight') + plt.close(fig) + buf.seek(0) + + pil_img = Image.open(buf).convert('RGB') + img_array = np.array(pil_img).astype(np.float32) / 255.0 + tensor_img = torch.from_numpy(img_array).unsqueeze(0) + + return tensor_img