Debugging tensors
This commit is contained in:
Pablo Apiolazza
+37
-148
@@ -1,5 +1,4 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
from ..utils.light_estimator import CategoricalLightEstimator
|
||||
from ..utils.luma_mask_processor import LumaMaskProcessor
|
||||
from ..utils.debug_visualizer import DebugVisualizer
|
||||
@@ -17,173 +16,63 @@ class NormalMapLightEstimator:
|
||||
"required": {
|
||||
"normal_map": ("IMAGE",),
|
||||
"luma_image": ("IMAGE",),
|
||||
"luma_threshold": ("FLOAT", {
|
||||
"default": 0.5,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01
|
||||
}),
|
||||
"luma_threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
"curve_type": (["linear", "s_curve", "exponential", "logarithmic"], {"default": "s_curve"}),
|
||||
# Directional thresholds
|
||||
"x_threshold": ("FLOAT", {
|
||||
"default": 0.1,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01
|
||||
}),
|
||||
"y_threshold": ("FLOAT", {
|
||||
"default": 0.1,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01
|
||||
}),
|
||||
"central_threshold": ("FLOAT", {
|
||||
"default": 0.3,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01
|
||||
}),
|
||||
# Hard/Soft classification thresholds
|
||||
"hard_light_threshold": ("FLOAT", {
|
||||
"default": 0.15,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01
|
||||
}),
|
||||
"soft_light_threshold": ("FLOAT", {
|
||||
"default": 0.35,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01
|
||||
}),
|
||||
"use_weighted": ("BOOLEAN", {"default": False}),
|
||||
"x_threshold": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
"y_threshold": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
"central_threshold": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
"hard_light_threshold": ("FLOAT", {"default": 0.15, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
"soft_light_threshold": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
},
|
||||
"optional": {
|
||||
"custom_curve_points": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": True
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = (
|
||||
"STRING", # x_direction
|
||||
"STRING", # y_direction
|
||||
"STRING", # combined_direction
|
||||
"FLOAT", # hard_soft_index
|
||||
"FLOAT", # x_confidence
|
||||
"FLOAT", # y_confidence
|
||||
"FLOAT", # overall_confidence
|
||||
"FLOAT", # spread_value
|
||||
"IMAGE", # debug_mask
|
||||
"IMAGE", # directional_visualization
|
||||
"STRING", "STRING", "STRING", "FLOAT", "FLOAT", "FLOAT", "FLOAT", "FLOAT",
|
||||
"IMAGE", "IMAGE", "IMAGE", "IMAGE",
|
||||
)
|
||||
|
||||
|
||||
RETURN_NAMES = (
|
||||
"x_direction",
|
||||
"y_direction",
|
||||
"combined_direction",
|
||||
"hard_soft_index",
|
||||
"x_confidence",
|
||||
"y_confidence",
|
||||
"overall_confidence",
|
||||
"spread_value",
|
||||
"debug_mask",
|
||||
"directional_viz"
|
||||
"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",
|
||||
)
|
||||
|
||||
FUNCTION = "estimate_lighting"
|
||||
|
||||
def estimate_lighting(self, normal_map, luma_image, luma_threshold, curve_type,
|
||||
x_threshold, y_threshold, central_threshold,
|
||||
hard_light_threshold, soft_light_threshold, use_weighted,
|
||||
custom_curve_points=""):
|
||||
hard_light_threshold, soft_light_threshold):
|
||||
"""
|
||||
Enhanced categorical light direction estimation with hard/soft index.
|
||||
Enhanced categorical light direction estimation.
|
||||
"""
|
||||
# Initialize processors
|
||||
luma_processor = LumaMaskProcessor()
|
||||
light_estimator = CategoricalLightEstimator()
|
||||
|
||||
# Set all thresholds
|
||||
light_estimator.x_threshold = x_threshold
|
||||
light_estimator.y_threshold = y_threshold
|
||||
light_estimator.central_threshold = central_threshold
|
||||
light_estimator.hard_light_threshold = hard_light_threshold
|
||||
light_estimator.soft_light_threshold = soft_light_threshold
|
||||
|
||||
# Process luma mask with curves
|
||||
if curve_type == "custom" and custom_curve_points:
|
||||
curve_points = self._parse_curve_points(custom_curve_points)
|
||||
mask = luma_processor.process_with_curves(
|
||||
luma_image, "custom", curve_points, luma_threshold
|
||||
)
|
||||
else:
|
||||
mask = luma_processor.process_with_curves(
|
||||
luma_image, curve_type, None, luma_threshold
|
||||
)
|
||||
|
||||
# Analyze directional categories with hard/soft index
|
||||
categorical_results = light_estimator.analyze_directional_categories(
|
||||
normal_map, mask
|
||||
light_estimator = CategoricalLightEstimator(
|
||||
x_threshold, y_threshold, central_threshold,
|
||||
hard_light_threshold, soft_light_threshold
|
||||
)
|
||||
|
||||
# Extract results for each batch
|
||||
x_directions = []
|
||||
y_directions = []
|
||||
combined_directions = []
|
||||
hard_soft_indices = []
|
||||
x_confidences = []
|
||||
y_confidences = []
|
||||
overall_confidences = []
|
||||
spread_values = []
|
||||
mask = luma_processor.process_with_curves(
|
||||
luma_image, curve_type, None, luma_threshold
|
||||
)
|
||||
|
||||
for result in categorical_results:
|
||||
x_directions.append(result['x_category'])
|
||||
y_directions.append(result['y_category'])
|
||||
combined_directions.append(result['combined_category'])
|
||||
hard_soft_indices.append(result['hard_soft_index'])
|
||||
x_confidences.append(result['confidence']['x_confidence'])
|
||||
y_confidences.append(result['confidence']['y_confidence'])
|
||||
overall_confidences.append(result['confidence']['overall_confidence'])
|
||||
spread_values.append(result['quality_analysis']['spread'])
|
||||
results = light_estimator.analyze_directional_categories(normal_map, mask)
|
||||
|
||||
x_direction = results['x_category']
|
||||
y_direction = results['y_category']
|
||||
combined_direction = f"{y_direction}-{x_direction}"
|
||||
hard_soft_index = results['hard_soft_index']
|
||||
x_confidence = results['confidence']['x_confidence']
|
||||
y_confidence = results['confidence']['y_confidence']
|
||||
overall_confidence = results['confidence']['overall_confidence']
|
||||
spread_value = results['quality_analysis']['spread']
|
||||
|
||||
# Generate debug outputs
|
||||
debug_mask = DebugVisualizer.generate_debug_mask(mask)
|
||||
directional_viz = DebugVisualizer.generate_directional_visualization(
|
||||
normal_map, categorical_results
|
||||
)
|
||||
|
||||
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()
|
||||
|
||||
return (
|
||||
x_directions,
|
||||
y_directions,
|
||||
combined_directions,
|
||||
torch.tensor(hard_soft_indices),
|
||||
torch.tensor(x_confidences),
|
||||
torch.tensor(y_confidences),
|
||||
torch.tensor(overall_confidences),
|
||||
torch.tensor(spread_values),
|
||||
debug_mask,
|
||||
directional_viz
|
||||
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
|
||||
)
|
||||
|
||||
def _parse_curve_points(self, curve_points_str):
|
||||
"""
|
||||
Parse custom curve points from string input.
|
||||
|
||||
Args:
|
||||
curve_points_str: String with curve points in format "x,y" per line
|
||||
|
||||
Returns:
|
||||
curve_points: List of (x, y) tuples
|
||||
"""
|
||||
try:
|
||||
points = []
|
||||
for line in curve_points_str.strip().split('\n'):
|
||||
if line.strip():
|
||||
x, y = map(float, line.strip().split(','))
|
||||
points.append((x, y))
|
||||
return points
|
||||
except:
|
||||
# Return default linear curve if parsing fails
|
||||
return [(0.0, 0.0), (1.0, 1.0)]
|
||||
|
||||
+69
-220
@@ -1,5 +1,4 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
class DebugVisualizer:
|
||||
"""
|
||||
@@ -10,234 +9,84 @@ class DebugVisualizer:
|
||||
def generate_debug_mask(mask):
|
||||
"""
|
||||
Generate debug visualization of the binary mask.
|
||||
|
||||
Args:
|
||||
mask: Binary mask tensor (B, H, W) or (B, H, W, 1)
|
||||
|
||||
Returns:
|
||||
debug_mask: RGB debug image (B, H, W, 3)
|
||||
"""
|
||||
# Handle different tensor shapes - squeeze extra dimensions if present
|
||||
# Ensure the mask is 3D (B, H, W) before unsqueezing
|
||||
if mask.dim() == 4 and mask.shape[-1] == 1:
|
||||
mask = mask.squeeze(-1)
|
||||
elif mask.dim() > 3:
|
||||
# If more than 3 dimensions, take first 3 and squeeze the rest
|
||||
mask = mask.flatten(0, mask.dim() - 3).squeeze()
|
||||
|
||||
# Add channel dimension and repeat for RGB
|
||||
debug_mask = mask.float().unsqueeze(-1).repeat(1, 1, 1, 3)
|
||||
return debug_mask
|
||||
|
||||
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_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_normal_visualization(normals, mask=None):
|
||||
def generate_colormap_preview(height=256, width=256):
|
||||
"""
|
||||
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)
|
||||
Generate a visual preview of the colormap used for analysis.
|
||||
"""
|
||||
# Convert normals from [-1, 1] to [0, 1] for RGB display
|
||||
normal_viz = (normals + 1.0) / 2.0
|
||||
x = torch.linspace(-1, 1, width)
|
||||
y = torch.linspace(-1, 1, height)
|
||||
xx, yy = torch.meshgrid(y, x, indexing='ij')
|
||||
|
||||
if mask is not None:
|
||||
# Dim unlit areas
|
||||
normal_viz = normal_viz * mask.unsqueeze(-1)
|
||||
# Mock results for visualization
|
||||
mock_normals = torch.stack([xx, yy, torch.zeros_like(xx)], dim=-1).unsqueeze(0)
|
||||
|
||||
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
|
||||
# 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]
|
||||
|
||||
# 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)
|
||||
# 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 torch.stack(visualizations)
|
||||
return colormap.byte().float() / 255.0
|
||||
|
||||
+80
-185
@@ -1,20 +1,20 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
class CategoricalLightEstimator:
|
||||
"""
|
||||
Light estimator with categorical directional outputs and simplified hard/soft classification.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
def __init__(self, x_threshold=0.1, y_threshold=0.1, central_threshold=0.3,
|
||||
hard_light_threshold=0.15, soft_light_threshold=0.35):
|
||||
# Directional thresholds
|
||||
self.x_threshold = 0.1
|
||||
self.y_threshold = 0.1
|
||||
self.central_threshold = 0.3
|
||||
self.x_threshold = x_threshold
|
||||
self.y_threshold = y_threshold
|
||||
self.central_threshold = central_threshold
|
||||
|
||||
# Hard/Soft classification thresholds
|
||||
self.hard_light_threshold = 0.15 # Below this = hard light
|
||||
self.soft_light_threshold = 0.35 # Above this = soft light
|
||||
self.hard_light_threshold = hard_light_threshold
|
||||
self.soft_light_threshold = soft_light_threshold
|
||||
|
||||
def analyze_directional_categories(self, normals, mask):
|
||||
"""
|
||||
@@ -25,252 +25,147 @@ class CategoricalLightEstimator:
|
||||
mask: Binary mask tensor (B, H, W)
|
||||
|
||||
Returns:
|
||||
results: List of analysis results for each batch
|
||||
results: Dictionary of analysis results for the batch
|
||||
"""
|
||||
batch_size = normals.shape[0]
|
||||
results = []
|
||||
batch_size, _, _, _ = normals.shape
|
||||
|
||||
for b in range(batch_size):
|
||||
# Get mask for current batch and ensure it's boolean
|
||||
current_mask = mask[b].bool()
|
||||
# Ensure mask is boolean
|
||||
current_mask = mask.bool()
|
||||
|
||||
# Apply mask to normals while preserving dimensions
|
||||
lit_normals = normals[b][current_mask]
|
||||
# Apply mask to normals
|
||||
lit_normals = normals[current_mask]
|
||||
|
||||
# Reshape to ensure 2D structure [N, 3] where N is number of lit pixels
|
||||
if lit_normals.dim() == 1:
|
||||
lit_normals = lit_normals.unsqueeze(0) if lit_normals.numel() > 0 else torch.empty(0, 3, device=normals.device, dtype=normals.dtype)
|
||||
if lit_normals.numel() == 0:
|
||||
return self._empty_categories(batch_size, normals.device)
|
||||
|
||||
if lit_normals.shape[0] == 0:
|
||||
results.append(self._empty_categories())
|
||||
continue
|
||||
if lit_normals.dim() == 1:
|
||||
lit_normals = lit_normals.unsqueeze(0)
|
||||
|
||||
# Extract XY components
|
||||
xy_normals = lit_normals[:, :2]
|
||||
|
||||
# Analyze directional distribution
|
||||
x_analysis = self._analyze_x_direction(xy_normals)
|
||||
y_analysis = self._analyze_y_direction(xy_normals)
|
||||
|
||||
# Analyze light quality (hard/soft index)
|
||||
quality_analysis = self._analyze_light_quality(xy_normals)
|
||||
|
||||
# Determine categories
|
||||
x_category = self._classify_x_direction(x_analysis)
|
||||
y_category = self._classify_y_direction(y_analysis)
|
||||
combined_category = f"{y_category}-{x_category}"
|
||||
|
||||
# Calculate hard/soft index (0 = hard, 1 = soft)
|
||||
hard_soft_index = self._calculate_hard_soft_index(quality_analysis)
|
||||
|
||||
# Calculate confidence scores
|
||||
confidence = {
|
||||
'x_confidence': x_analysis['confidence'],
|
||||
'y_confidence': y_analysis['confidence'],
|
||||
'quality_confidence': quality_analysis['confidence'],
|
||||
'overall_confidence': (x_analysis['confidence'] + y_analysis['confidence'] + quality_analysis['confidence']) / 3
|
||||
}
|
||||
|
||||
results.append({
|
||||
'x_category': x_category,
|
||||
'y_category': y_category,
|
||||
'combined_category': combined_category,
|
||||
'hard_soft_index': hard_soft_index,
|
||||
'confidence': confidence,
|
||||
'x_analysis': x_analysis,
|
||||
'y_analysis': y_analysis,
|
||||
'quality_analysis': quality_analysis
|
||||
})
|
||||
# Extract XY components
|
||||
xy_normals = lit_normals[:, :2]
|
||||
|
||||
return results
|
||||
# Analyze directional distribution
|
||||
x_analysis = self._analyze_x_direction(xy_normals)
|
||||
y_analysis = self._analyze_y_direction(xy_normals)
|
||||
|
||||
# Analyze light quality
|
||||
quality_analysis = self._analyze_light_quality(xy_normals)
|
||||
|
||||
# Determine categories
|
||||
x_category = self._classify_x_direction(x_analysis)
|
||||
y_category = self._classify_y_direction(y_analysis)
|
||||
|
||||
# Calculate hard/soft index
|
||||
hard_soft_index = self._calculate_hard_soft_index(quality_analysis)
|
||||
|
||||
# Calculate confidence scores
|
||||
confidence = {
|
||||
'x_confidence': x_analysis['confidence'],
|
||||
'y_confidence': y_analysis['confidence'],
|
||||
'quality_confidence': quality_analysis['confidence'],
|
||||
'overall_confidence': (x_analysis['confidence'] + y_analysis['confidence'] + quality_analysis['confidence']) / 3
|
||||
}
|
||||
|
||||
return {
|
||||
'x_category': x_category,
|
||||
'y_category': y_category,
|
||||
'hard_soft_index': hard_soft_index,
|
||||
'confidence': confidence,
|
||||
'quality_analysis': quality_analysis
|
||||
}
|
||||
|
||||
def _analyze_light_quality(self, xy_normals):
|
||||
"""
|
||||
Analyze light quality based on spread of normals.
|
||||
|
||||
Args:
|
||||
xy_normals: XY components of normal vectors (N, 2)
|
||||
|
||||
Returns:
|
||||
quality_analysis: Dict with spread and confidence
|
||||
"""
|
||||
if len(xy_normals) < 2:
|
||||
return {
|
||||
'spread': 0.0,
|
||||
'confidence': 0.0
|
||||
}
|
||||
if xy_normals.shape[0] < 2:
|
||||
return {'spread': 0.0, 'confidence': 0.0}
|
||||
|
||||
# Calculate spread using covariance analysis
|
||||
xy_np = xy_normals.detach().cpu().numpy()
|
||||
cov_matrix = np.cov(xy_np.T)
|
||||
eigenvalues = np.linalg.eigvals(cov_matrix)
|
||||
spread = np.sqrt(np.max(eigenvalues))
|
||||
# Calculate spread using covariance matrix
|
||||
cov_matrix = torch.cov(xy_normals.T)
|
||||
eigenvalues = torch.linalg.eigvals(cov_matrix).real
|
||||
spread = torch.sqrt(torch.max(eigenvalues))
|
||||
|
||||
# Calculate confidence based on how clear the spread is
|
||||
confidence = min(1.0, spread / 0.5) # Normalize confidence
|
||||
# Normalize confidence
|
||||
confidence = torch.clamp(spread / 0.5, 0.0, 1.0)
|
||||
|
||||
return {
|
||||
'spread': spread,
|
||||
'confidence': confidence
|
||||
}
|
||||
return {'spread': spread.item(), 'confidence': confidence.item()}
|
||||
|
||||
def _calculate_hard_soft_index(self, quality_analysis):
|
||||
"""
|
||||
Calculate hard/soft index from 0 (hard) to 1 (soft).
|
||||
|
||||
Args:
|
||||
quality_analysis: Analysis results with spread value
|
||||
|
||||
Returns:
|
||||
hard_soft_index: Float from 0.0 (hard) to 1.0 (soft)
|
||||
"""
|
||||
spread = quality_analysis['spread']
|
||||
|
||||
# Linear interpolation between thresholds
|
||||
if spread <= self.hard_light_threshold:
|
||||
# Hard light range
|
||||
return 0.0
|
||||
elif spread >= self.soft_light_threshold:
|
||||
# Soft light range
|
||||
return 1.0
|
||||
else:
|
||||
# Intermediate range - linear interpolation
|
||||
ratio = (spread - self.hard_light_threshold) / (self.soft_light_threshold - self.hard_light_threshold)
|
||||
return ratio
|
||||
return (spread - self.hard_light_threshold) / (self.soft_light_threshold - self.hard_light_threshold)
|
||||
|
||||
def _analyze_x_direction(self, xy_normals):
|
||||
"""
|
||||
Analyze X-direction distribution and calculate confidence.
|
||||
|
||||
Args:
|
||||
xy_normals: XY components of normal vectors (N, 2)
|
||||
|
||||
Returns:
|
||||
x_analysis: Dict with percentages and confidence
|
||||
Analyze X-direction distribution.
|
||||
"""
|
||||
x_components = xy_normals[:, 0]
|
||||
total_count = x_components.shape[0]
|
||||
|
||||
# Count normals in each X category
|
||||
left_count = torch.sum(x_components < -self.x_threshold).item()
|
||||
right_count = torch.sum(x_components > self.x_threshold).item()
|
||||
central_count = torch.sum(torch.abs(x_components) <= self.central_threshold).item()
|
||||
left_pct = torch.sum(x_components < -self.x_threshold).item() / total_count
|
||||
right_pct = torch.sum(x_components > self.x_threshold).item() / total_count
|
||||
central_pct = torch.sum(torch.abs(x_components) <= self.central_threshold).item() / total_count
|
||||
|
||||
total_count = len(x_components)
|
||||
|
||||
# Calculate percentages
|
||||
left_pct = left_count / total_count if total_count > 0 else 0
|
||||
right_pct = right_count / total_count if total_count > 0 else 0
|
||||
central_pct = central_count / total_count if total_count > 0 else 0
|
||||
|
||||
# Calculate confidence
|
||||
max_pct = max(left_pct, right_pct, central_pct)
|
||||
confidence = max_pct if max_pct > 0.3 else 0.0
|
||||
|
||||
return {
|
||||
'left_pct': left_pct,
|
||||
'right_pct': right_pct,
|
||||
'central_pct': central_pct,
|
||||
'confidence': confidence,
|
||||
'total_normals': total_count
|
||||
}
|
||||
return {'left_pct': left_pct, 'right_pct': right_pct, 'central_pct': central_pct, 'confidence': confidence}
|
||||
|
||||
def _analyze_y_direction(self, xy_normals):
|
||||
"""
|
||||
Analyze Y-direction distribution and calculate confidence.
|
||||
|
||||
Args:
|
||||
xy_normals: XY components of normal vectors (N, 2)
|
||||
|
||||
Returns:
|
||||
y_analysis: Dict with percentages and confidence
|
||||
Analyze Y-direction distribution.
|
||||
"""
|
||||
y_components = xy_normals[:, 1]
|
||||
total_count = y_components.shape[0]
|
||||
|
||||
top_pct = torch.sum(y_components > self.y_threshold).item() / total_count
|
||||
bottom_pct = torch.sum(y_components < -self.y_threshold).item() / total_count
|
||||
central_pct = torch.sum(torch.abs(y_components) <= self.central_threshold).item() / total_count
|
||||
|
||||
# Count normals in each Y category
|
||||
top_count = torch.sum(y_components > self.y_threshold).item()
|
||||
bottom_count = torch.sum(y_components < -self.y_threshold).item()
|
||||
central_count = torch.sum(torch.abs(y_components) <= self.central_threshold).item()
|
||||
|
||||
total_count = len(y_components)
|
||||
|
||||
# Calculate percentages
|
||||
top_pct = top_count / total_count if total_count > 0 else 0
|
||||
bottom_pct = bottom_count / total_count if total_count > 0 else 0
|
||||
central_pct = central_count / total_count if total_count > 0 else 0
|
||||
|
||||
# Calculate confidence
|
||||
max_pct = max(top_pct, bottom_pct, central_pct)
|
||||
confidence = max_pct if max_pct > 0.3 else 0.0
|
||||
|
||||
return {
|
||||
'top_pct': top_pct,
|
||||
'bottom_pct': bottom_pct,
|
||||
'central_pct': central_pct,
|
||||
'confidence': confidence,
|
||||
'total_normals': total_count
|
||||
}
|
||||
return {'top_pct': top_pct, 'bottom_pct': bottom_pct, 'central_pct': central_pct, 'confidence': confidence}
|
||||
|
||||
def _classify_x_direction(self, x_analysis):
|
||||
"""
|
||||
Classify X direction based on analysis.
|
||||
|
||||
Args:
|
||||
x_analysis: X direction analysis results
|
||||
|
||||
Returns:
|
||||
x_category: "left", "central", or "right"
|
||||
Classify X direction.
|
||||
"""
|
||||
left_pct = x_analysis['left_pct']
|
||||
right_pct = x_analysis['right_pct']
|
||||
central_pct = x_analysis['central_pct']
|
||||
|
||||
if central_pct > max(left_pct, right_pct):
|
||||
if x_analysis['central_pct'] > max(x_analysis['left_pct'], x_analysis['right_pct']):
|
||||
return "central"
|
||||
elif left_pct > right_pct:
|
||||
elif x_analysis['left_pct'] > x_analysis['right_pct']:
|
||||
return "left"
|
||||
else:
|
||||
return "right"
|
||||
|
||||
def _classify_y_direction(self, y_analysis):
|
||||
"""
|
||||
Classify Y direction based on analysis.
|
||||
|
||||
Args:
|
||||
y_analysis: Y direction analysis results
|
||||
|
||||
Returns:
|
||||
y_category: "top", "central", or "bottom"
|
||||
Classify Y direction.
|
||||
"""
|
||||
top_pct = y_analysis['top_pct']
|
||||
bottom_pct = y_analysis['bottom_pct']
|
||||
central_pct = y_analysis['central_pct']
|
||||
|
||||
if central_pct > max(top_pct, bottom_pct):
|
||||
if y_analysis['central_pct'] > max(y_analysis['top_pct'], y_analysis['bottom_pct']):
|
||||
return "central"
|
||||
elif top_pct > bottom_pct:
|
||||
elif y_analysis['top_pct'] > y_analysis['bottom_pct']:
|
||||
return "top"
|
||||
else:
|
||||
return "bottom"
|
||||
|
||||
def _empty_categories(self):
|
||||
def _empty_categories(self, batch_size, device):
|
||||
"""
|
||||
Return empty categories for cases with no lit normals.
|
||||
|
||||
Returns:
|
||||
empty_result: Dict with default values
|
||||
"""
|
||||
return {
|
||||
'x_category': "central",
|
||||
'y_category': "central",
|
||||
'combined_category': "central-central",
|
||||
'hard_soft_index': 0.5, # Default to intermediate
|
||||
'confidence': {
|
||||
'x_confidence': 0.0,
|
||||
'y_confidence': 0.0,
|
||||
'quality_confidence': 0.0,
|
||||
'overall_confidence': 0.0
|
||||
},
|
||||
'x_analysis': {'left_pct': 0.0, 'right_pct': 0.0, 'central_pct': 0.0, 'confidence': 0.0, 'total_normals': 0},
|
||||
'y_analysis': {'top_pct': 0.0, 'bottom_pct': 0.0, 'central_pct': 0.0, 'confidence': 0.0, 'total_normals': 0},
|
||||
'hard_soft_index': 0.5,
|
||||
'confidence': {'x_confidence': 0.0, 'y_confidence': 0.0, 'quality_confidence': 0.0, 'overall_confidence': 0.0},
|
||||
'quality_analysis': {'spread': 0.0, 'confidence': 0.0}
|
||||
}
|
||||
|
||||
@@ -21,9 +21,13 @@ class LumaMaskProcessor:
|
||||
Returns:
|
||||
mask: Binary mask tensor (B, H, W) where True indicates lit areas
|
||||
"""
|
||||
# Ensure grayscale by averaging channels if necessary
|
||||
if luma_image.shape[-1] > 1:
|
||||
luma_image = luma_image.mean(dim=-1, keepdim=True)
|
||||
|
||||
# Convert to numpy for curve processing
|
||||
luma_np = luma_image.detach().cpu().numpy()
|
||||
|
||||
|
||||
if curve_type == "linear":
|
||||
# Simple linear mapping
|
||||
processed_luma = luma_np
|
||||
|
||||
Reference in New Issue
Block a user