diff --git a/nodes/nodes.py b/nodes/nodes.py index 0e5e12f..0887d23 100644 --- a/nodes/nodes.py +++ b/nodes/nodes.py @@ -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)] diff --git a/utils/debug_visualizer.py b/utils/debug_visualizer.py index 8938691..9bbb854 100644 --- a/utils/debug_visualizer.py +++ b/utils/debug_visualizer.py @@ -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 diff --git a/utils/light_estimator.py b/utils/light_estimator.py index 63a3988..2dedca3 100644 --- a/utils/light_estimator.py +++ b/utils/light_estimator.py @@ -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} } diff --git a/utils/luma_mask_processor.py b/utils/luma_mask_processor.py index 28c8f47..eb9606c 100644 --- a/utils/luma_mask_processor.py +++ b/utils/luma_mask_processor.py @@ -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