409 lines
15 KiB
Python
409 lines
15 KiB
Python
import torch
|
|
|
|
class DebugVisualizer:
|
|
"""
|
|
Debug visualization tools for light estimation analysis.
|
|
"""
|
|
|
|
@staticmethod
|
|
def generate_debug_mask(mask):
|
|
"""
|
|
Generate debug visualization of the binary mask.
|
|
"""
|
|
# Ensure mask is 3D (B, H, W)
|
|
if mask.dim() == 4 and mask.shape[-1] == 1:
|
|
mask = mask.squeeze(-1)
|
|
|
|
# Convert to RGB visualization: white for lit areas, black for unlit
|
|
debug_mask = mask.float().unsqueeze(-1).repeat(1, 1, 1, 3)
|
|
return debug_mask
|
|
|
|
@staticmethod
|
|
def generate_lit_normals_visualization(normals, mask):
|
|
"""
|
|
Generate visualization showing only the lit normal vectors.
|
|
"""
|
|
normal_rgb = (normals + 1.0) / 2.0
|
|
lit_mask = mask.unsqueeze(-1)
|
|
lit_normals_viz = normal_rgb * lit_mask
|
|
return lit_normals_viz
|
|
|
|
|
|
@staticmethod
|
|
def generate_colormap_preview(height=512, width=512):
|
|
"""
|
|
Generate a comprehensive reference guide for normal map color encoding.
|
|
|
|
This creates an educational reference showing exactly how 3D surface
|
|
orientations are encoded as RGB colors in normal maps.
|
|
"""
|
|
# Create main visualization canvas
|
|
viz = torch.ones(height, width, 3) * 0.9 # Light gray background
|
|
|
|
# Title section
|
|
title_height = 60
|
|
viz[:title_height, :] = torch.tensor([0.7, 0.7, 0.7]) # Darker background for title
|
|
|
|
# Color reference sections in a 2x3 grid
|
|
section_height = (height - title_height - 40) // 2 # Leave space for labels
|
|
section_width = width // 3
|
|
|
|
# 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 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 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
|
|
|
|
viz[y_center:y_center+color_block_size, x_center:x_center+color_block_size] = direction['color']
|
|
|
|
# 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)
|
|
|
|
@staticmethod
|
|
def generate_color_histogram(normals, title="Normal Distribution"):
|
|
"""
|
|
Generate a clear, informative histogram of normal map color distribution.
|
|
|
|
Args:
|
|
normals: Normal vectors tensor (B, H, W, 3) in [-1, 1] range
|
|
title: Title for the histogram
|
|
|
|
Returns:
|
|
histogram_viz: Clear visualization with labels and reference (1, 512, 512, 3)
|
|
"""
|
|
return DebugVisualizer._create_enhanced_histogram(normals, title, show_reference=True)
|
|
|
|
@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)
|
|
|
|
# Sample for performance
|
|
sample_size = min(flat_normals.shape[0], 100000)
|
|
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()
|
|
|
|
# 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 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"""
|
|
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
|