237 lines
8.1 KiB
Python
237 lines
8.1 KiB
Python
import torch
|
|
import numpy as np
|
|
|
|
class DebugVisualizer:
|
|
"""
|
|
Debug visualization tools for light estimation analysis.
|
|
"""
|
|
|
|
@staticmethod
|
|
def generate_debug_mask(mask):
|
|
"""
|
|
Generate debug visualization of the binary mask.
|
|
|
|
Args:
|
|
mask: Binary mask tensor (B, H, W)
|
|
|
|
Returns:
|
|
debug_mask: RGB debug image (B, H, W, 3)
|
|
"""
|
|
batch_size, height, width = mask.shape
|
|
debug_masks = []
|
|
|
|
for b in range(batch_size):
|
|
# Convert binary mask to RGB
|
|
mask_np = mask[b].detach().cpu().numpy()
|
|
debug_rgb = np.zeros((height, width, 3), dtype=np.uint8)
|
|
|
|
# Set lit areas to white, unlit to black
|
|
debug_rgb[mask_np] = [255, 255, 255]
|
|
|
|
debug_masks.append(torch.from_numpy(debug_rgb).float() / 255.0)
|
|
|
|
return torch.stack(debug_masks)
|
|
|
|
@staticmethod
|
|
def generate_normal_visualization(normals, mask=None):
|
|
"""
|
|
Generate RGB visualization of normal vectors.
|
|
|
|
Args:
|
|
normals: Normal vectors (B, H, W, 3)
|
|
mask: Optional binary mask to highlight specific areas
|
|
|
|
Returns:
|
|
normal_viz: RGB visualization (B, H, W, 3)
|
|
"""
|
|
# Convert normals from [-1, 1] to [0, 1] for RGB display
|
|
normal_viz = (normals + 1.0) / 2.0
|
|
|
|
if mask is not None:
|
|
# Dim unlit areas
|
|
normal_viz = normal_viz * mask.unsqueeze(-1)
|
|
|
|
return normal_viz
|
|
|
|
@staticmethod
|
|
def generate_spread_histogram(normals, mask, bins=36):
|
|
"""
|
|
Generate polar histogram of lit normal distribution.
|
|
|
|
Args:
|
|
normals: Normal vectors (B, H, W, 3)
|
|
mask: Binary mask (B, H, W)
|
|
bins: Number of histogram bins
|
|
|
|
Returns:
|
|
histogram: Polar histogram data
|
|
"""
|
|
batch_size = normals.shape[0]
|
|
histograms = []
|
|
|
|
for b in range(batch_size):
|
|
lit_normals = normals[b][mask[b]]
|
|
|
|
if len(lit_normals) == 0:
|
|
histograms.append(np.zeros(bins))
|
|
continue
|
|
|
|
# Extract XY components and convert to polar coordinates
|
|
xy_normals = lit_normals[:, :2]
|
|
angles = torch.atan2(xy_normals[:, 1], xy_normals[:, 0])
|
|
|
|
# Convert to degrees and create histogram
|
|
angles_deg = torch.rad2deg(angles).detach().cpu().numpy()
|
|
hist, _ = np.histogram(angles_deg, bins=bins, range=(-180, 180))
|
|
histograms.append(hist)
|
|
|
|
return histograms
|
|
|
|
@staticmethod
|
|
def generate_directional_visualization(normals, categorical_results):
|
|
"""
|
|
Generate visualization showing categorical analysis results.
|
|
|
|
Args:
|
|
normals: Normal vectors (B, H, W, 3)
|
|
categorical_results: Results from light estimator
|
|
|
|
Returns:
|
|
directional_viz: Color-coded visualization (B, H, W, 3)
|
|
"""
|
|
batch_size = normals.shape[0]
|
|
visualizations = []
|
|
|
|
for b in range(batch_size):
|
|
result = categorical_results[b]
|
|
|
|
# Create color-coded visualization
|
|
height, width = normals.shape[1], normals.shape[2]
|
|
viz = np.zeros((height, width, 3), dtype=np.uint8)
|
|
|
|
# Get categories and index
|
|
x_category = result['x_category']
|
|
y_category = result['y_category']
|
|
hard_soft_index = result['hard_soft_index']
|
|
overall_confidence = result['confidence']['overall_confidence']
|
|
|
|
# X direction colors (Red channel)
|
|
if x_category == "left":
|
|
red = 255
|
|
elif x_category == "right":
|
|
red = 128
|
|
else: # central
|
|
red = 192
|
|
|
|
# Y direction colors (Green channel)
|
|
if y_category == "top":
|
|
green = 255
|
|
elif y_category == "bottom":
|
|
green = 128
|
|
else: # central
|
|
green = 192
|
|
|
|
# Hard/Soft index colors (Blue channel)
|
|
# 0.0 (hard) = 255 (bright blue)
|
|
# 1.0 (soft) = 128 (medium blue)
|
|
blue = int(255 - (hard_soft_index * 127))
|
|
|
|
# Apply confidence as alpha/transparency effect
|
|
confidence_factor = overall_confidence
|
|
red = int(red * confidence_factor)
|
|
green = int(green * confidence_factor)
|
|
blue = int(blue * confidence_factor)
|
|
|
|
viz[:, :] = [red, green, blue]
|
|
visualizations.append(torch.from_numpy(viz).float() / 255.0)
|
|
|
|
return torch.stack(visualizations)
|
|
|
|
@staticmethod
|
|
def generate_confidence_heatmap(normals, categorical_results):
|
|
"""
|
|
Generate confidence heatmap visualization.
|
|
|
|
Args:
|
|
normals: Normal vectors (B, H, W, 3)
|
|
categorical_results: Results from light estimator
|
|
|
|
Returns:
|
|
confidence_viz: Confidence heatmap (B, H, W, 3)
|
|
"""
|
|
batch_size = normals.shape[0]
|
|
visualizations = []
|
|
|
|
for b in range(batch_size):
|
|
result = categorical_results[b]
|
|
|
|
# Create confidence heatmap
|
|
height, width = normals.shape[1], normals.shape[2]
|
|
viz = np.zeros((height, width, 3), dtype=np.uint8)
|
|
|
|
# Get confidence values
|
|
x_confidence = result['confidence']['x_confidence']
|
|
y_confidence = result['confidence']['y_confidence']
|
|
quality_confidence = result['confidence']['quality_confidence']
|
|
overall_confidence = result['confidence']['overall_confidence']
|
|
|
|
# Color code based on confidence levels
|
|
# Red = X confidence, Green = Y confidence, Blue = Quality confidence
|
|
red = int(255 * x_confidence)
|
|
green = int(255 * y_confidence)
|
|
blue = int(255 * quality_confidence)
|
|
|
|
viz[:, :] = [red, green, blue]
|
|
visualizations.append(torch.from_numpy(viz).float() / 255.0)
|
|
|
|
return torch.stack(visualizations)
|
|
|
|
@staticmethod
|
|
def generate_quadrant_visualization(normals, mask):
|
|
"""
|
|
Generate visualization showing directional quadrants.
|
|
|
|
Args:
|
|
normals: Normal vectors (B, H, W, 3)
|
|
mask: Binary mask (B, H, W)
|
|
|
|
Returns:
|
|
quadrant_viz: Quadrant visualization (B, H, W, 3)
|
|
"""
|
|
batch_size = normals.shape[0]
|
|
visualizations = []
|
|
|
|
for b in range(batch_size):
|
|
# Create quadrant visualization
|
|
height, width = normals.shape[1], normals.shape[2]
|
|
viz = np.zeros((height, width, 3), dtype=np.uint8)
|
|
|
|
# Get lit normals
|
|
lit_normals = normals[b][mask[b]]
|
|
|
|
if len(lit_normals) == 0:
|
|
visualizations.append(torch.from_numpy(viz).float() / 255.0)
|
|
continue
|
|
|
|
# Extract XY components
|
|
xy_normals = lit_normals[:, :2]
|
|
|
|
# Color code based on quadrants
|
|
for i, (x, y) in enumerate(xy_normals):
|
|
if x > 0 and y > 0: # Top-right
|
|
color = [255, 255, 0] # Yellow
|
|
elif x < 0 and y > 0: # Top-left
|
|
color = [255, 0, 255] # Magenta
|
|
elif x < 0 and y < 0: # Bottom-left
|
|
color = [0, 255, 255] # Cyan
|
|
else: # Bottom-right
|
|
color = [255, 0, 0] # Red
|
|
|
|
# Find pixel position (simplified - would need proper mapping)
|
|
# This is a placeholder for the actual implementation
|
|
pass
|
|
|
|
visualizations.append(torch.from_numpy(viz).float() / 255.0)
|
|
|
|
return torch.stack(visualizations)
|