124 lines
4.5 KiB
Python
124 lines
4.5 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_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):
|
|
"""
|
|
Generate a visual preview of the colormap used for analysis.
|
|
"""
|
|
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
|
|
|
|
@staticmethod
|
|
def generate_color_histogram(normals, title="Color Distribution"):
|
|
"""
|
|
Generate a color histogram visualization of normal vectors.
|
|
|
|
Args:
|
|
normals: Normal vectors tensor (B, H, W, 3)
|
|
title: Title for the histogram
|
|
|
|
Returns:
|
|
histogram_viz: RGB visualization of color distribution (1, 256, 256, 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)
|
|
|
|
# Create histogram visualization
|
|
hist_size = 256
|
|
histogram = torch.zeros(hist_size, hist_size, 3)
|
|
|
|
# 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])
|
|
|
|
return histogram.unsqueeze(0).byte().float() / 255.0
|