commit 60adfbde322f4474dcd5d749325cfbff6adc3fc0 Author: Pablo Apiolazza <31888952+APZmedia@users.noreply.github.com> Date: Wed Oct 1 12:43:15 2025 +0200 First commit - POC diff --git a/DESIGN_NOTES.md b/DESIGN_NOTES.md new file mode 100644 index 0000000..3d2fe0c --- /dev/null +++ b/DESIGN_NOTES.md @@ -0,0 +1,192 @@ +# NormalMapLightEstimator - Design Notes + +## 🎯 Project Overview + +A ComfyUI custom node for estimating light direction and quality from normal maps using luma masking. The system analyzes surface normals to infer lighting information for downstream tasks like adaptive relighting, directional masking, or stylized effects. + +## 🧠 Core Algorithm + +### Physics-Based Approach +- **Normal Map (RGB)** → Surface orientation (XYZ) +- **Luma Mask** → Which areas are actually lit +- **Filtered Normals** → Only consider lit surface orientations +- **Statistical Analysis** → Infer light direction and quality + +### Key Insight +Lit surface normals point **TOWARD** the light source. By averaging these directions, we get the light direction. + +## 📦 Functional Requirements + +### Node Name +`NormalMapLightEstimator` + +### Inputs +| Name | Type | Description | +|------|------|-------------| +| `normal_map` | IMAGE | RGB normal map in tangent space (0–255) | +| `luma_image` | IMAGE | Grayscale or RGB image to extract luminance | +| `luma_threshold` | FLOAT | Threshold to mask out non-lit areas (default 0.5) | +| `curve_type` | COMBO | Luma processing curve: linear, s_curve, exponential, logarithmic | +| `x_threshold` | FLOAT | Sensitivity for left/right detection (default 0.1) | +| `y_threshold` | FLOAT | Sensitivity for top/bottom detection (default 0.1) | +| `central_threshold` | FLOAT | Threshold for central lighting detection (default 0.3) | +| `hard_light_threshold` | FLOAT | Below this = hard light (default 0.15) | +| `soft_light_threshold` | FLOAT | Above this = soft light (default 0.35) | +| `use_weighted` | BOOLEAN | Use weighted averaging instead of binary masking | +| `custom_curve_points` | STRING | Custom curve points for advanced luma processing | + +### Outputs +| Name | Type | Description | +|------|------|-------------| +| `x_direction` | STRING | "left", "central", "right" | +| `y_direction` | STRING | "top", "central", "bottom" | +| `combined_direction` | STRING | "top-left", "top-central", "top-right", etc. | +| `hard_soft_index` | FLOAT | 0.0 (hard) to 1.0 (soft) | +| `x_confidence` | FLOAT | Confidence in X direction classification | +| `y_confidence` | FLOAT | Confidence in Y direction classification | +| `overall_confidence` | FLOAT | Combined confidence score | +| `spread_value` | FLOAT | Raw numerical spread value | +| `debug_mask` | IMAGE | Binary mask showing which normals were used | +| `directional_viz` | IMAGE | Color-coded visualization of results | + +## ⚙️ Processing Pipeline + +### 1. Normal Map Processing +- Convert RGB to XYZ: `normals = ((RGB / 255) * 2.0) - 1.0` +- Apply axis inversions for different normal conventions +- Support for OpenGL vs DirectX normal maps + +### 2. Luma Mask Processing with Curves Strategy +- **Linear**: Simple threshold-like behavior +- **S-curve**: Enhance mid-tones, compress highlights/shadows +- **Exponential**: Emphasize bright areas +- **Logarithmic**: Compress highlights, expand shadows +- **Custom**: User-defined curve points + +### 3. Directional Analysis +- **X Direction**: Count normals in left/right/central quadrants +- **Y Direction**: Count normals in top/bottom/central quadrants +- **Central Detection**: Identify omnidirectional lighting patterns + +### 4. Hard/Soft Classification +- **Spread Analysis**: Use covariance analysis of lit normals +- **Hard/Soft Index**: Continuous value from 0.0 (hard) to 1.0 (soft) +- **Threshold Control**: User-adjustable hard/soft boundaries + +## 🎛️ Threshold Controls + +### Directional Thresholds +- **X Threshold**: Controls left/right detection sensitivity +- **Y Threshold**: Controls top/bottom detection sensitivity +- **Central Threshold**: Detects central lighting patterns + +### Quality Thresholds +- **Hard Light Threshold**: Below this = hard light (index 0.0) +- **Soft Light Threshold**: Above this = soft light (index 1.0) +- **Between**: Linear interpolation from 0.0 to 1.0 + +## 🎯 Output Categories + +### X Direction +- **"left"**: Light coming from the left +- **"central"**: Central/omnidirectional lighting +- **"right"**: Light coming from the right + +### Y Direction +- **"top"**: Light coming from above +- **"central"**: Central/omnidirectional lighting +- **"bottom"**: Light coming from below + +### Combined Directions +- **"top-left"**, **"top-central"**, **"top-right"** +- **"central-left"**, **"central-central"**, **"central-right"** +- **"bottom-left"**, **"bottom-central"**, **"bottom-right"** + +### Hard/Soft Index +- **0.0**: Hard light (tight cluster of normals) +- **0.5**: Intermediate light +- **1.0**: Soft light (wide distribution of normals) + +## 🔧 Technical Implementation + +### Modular Architecture +- **`normal_map_processor.py`**: Normal map processing (leverage existing ComfyUI nodes) +- **`luma_mask_processor.py`**: Luma processing with curves strategy +- **`light_estimator.py`**: Core light estimation algorithm +- **`debug_visualizer.py`**: Debug visualization tools +- **`nodes.py`**: Main ComfyUI node class + +### Key Algorithms +1. **Directional Analysis**: Quadrant-based normal distribution analysis +2. **Spread Calculation**: Covariance analysis of lit normals +3. **Confidence Scoring**: Statistical confidence in classifications +4. **Curves Processing**: Advanced luma masking with user control + +## 📊 Usage Examples + +### Basic Directional Detection +``` +X Threshold: 0.1 (sensitive to left/right) +Y Threshold: 0.1 (sensitive to top/bottom) +Central Threshold: 0.3 (detect central lighting) +``` + +### Hard Light Detection +``` +Hard Threshold: 0.1 (very sensitive) +Soft Threshold: 0.2 (narrow range) +Result: Index 0.0-0.5 range +``` + +### Soft Light Detection +``` +Hard Threshold: 0.2 (less sensitive) +Soft Threshold: 0.5 (wider range) +Result: Index 0.5-1.0 range +``` + +## 🎨 Visualization System + +### Debug Mask +- White pixels: Lit areas used in analysis +- Black pixels: Masked out areas + +### Directional Visualization +- **Red Channel**: X direction (left=255, central=192, right=128) +- **Green Channel**: Y direction (top=255, central=192, bottom=128) +- **Blue Channel**: Hard/soft index (hard=255, soft=128) +- **Confidence**: Applied as transparency effect + +## 🚀 Future Enhancements + +### Phase 2 Features +- Weighted averaging using luma intensity +- Multi-scale processing for different detail levels +- Polar histogram output for normal distribution +- Integration with existing ComfyUI lighting nodes + +### Phase 3 Features +- Real-time threshold adjustment +- Batch processing optimization +- Advanced curve presets +- Export/import threshold configurations + +## 📝 Implementation Notes + +### ComfyUI Integration +- **Category**: "Custom/Lighting" +- **Dependencies**: torch, numpy, Pillow +- **Auto-installation**: Dependencies install automatically +- **Error Handling**: Graceful fallbacks for edge cases + +### Performance Considerations +- Optimized for ComfyUI tensor operations +- Batch processing support +- Memory-efficient covariance calculations +- GPU acceleration where possible + +### Edge Cases Handled +- No lit areas: Returns default central direction +- Single normal: Uses that direction +- Collinear normals: Handles degenerate cases gracefully +- Low confidence: Provides reliability indicators diff --git a/INSTALL.md b/INSTALL.md new file mode 100644 index 0000000..628f108 --- /dev/null +++ b/INSTALL.md @@ -0,0 +1,68 @@ +# Installation Guide + +## Quick Installation + +1. **Copy to ComfyUI custom nodes directory:** + ``` + ComfyUI/custom_nodes/Comfyui-LightDirection-estimation/ + ``` + +2. **Restart ComfyUI** + +3. **Find the node** under "Custom/Lighting" category + +## Manual Installation + +If auto-installation fails: + +```bash +cd ComfyUI/custom_nodes/Comfyui-LightDirection-estimation/ +pip install -r requirements.txt +``` + +## Development Installation + +For development purposes: + +```bash +git clone https://github.com/your-username/Comfyui-LightDirection-estimation.git +cd Comfyui-LightDirection-estimation +pip install -e . +``` + +## Dependencies + +- torch>=1.9.0 +- numpy>=1.21.0 +- Pillow>=8.0.0 +- scipy>=1.7.0 + +## Troubleshooting + +### Common Issues + +**"Module not found" errors** +- Ensure all dependencies are installed +- Check Python path includes the custom_nodes directory + +**"Node not appearing in ComfyUI"** +- Verify the folder structure is correct +- Check that `__init__.py` files are present +- Restart ComfyUI completely + +**"Import errors"** +- Check that all required packages are installed +- Verify Python version compatibility (3.8+) + +### Manual Dependency Installation + +```bash +pip install torch>=1.9.0 numpy>=1.21.0 Pillow>=8.0.0 scipy>=1.7.0 +``` + +## Verification + +After installation, you should see: +- "Normal Map Light Estimator" node in ComfyUI under "Custom/Lighting" +- No import errors in ComfyUI console +- All dependencies properly installed diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..b6c5757 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 ComfyUI Light Direction Estimation + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000..b4bb48e --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1,6 @@ +include README.md +include requirements.txt +include LICENSE +recursive-include nodes *.py +recursive-include utils *.py +recursive-include example *.json *.md diff --git a/README.md b/README.md new file mode 100644 index 0000000..2b96a5e --- /dev/null +++ b/README.md @@ -0,0 +1,211 @@ +# NormalMapLightEstimator + +A ComfyUI custom node for estimating light direction and quality from normal maps using luma masking. The system analyzes surface normals to infer lighting information for downstream tasks like adaptive relighting, directional masking, or stylized effects. + +## Features + +- **Categorical Directional Outputs**: X direction (left/central/right), Y direction (top/central/bottom) +- **Hard/Soft Classification**: Continuous index from 0.0 (hard) to 1.0 (soft) +- **Curves Strategy**: Advanced luma processing with multiple curve types +- **User-Controllable Thresholds**: Fine-tune detection sensitivity +- **Debug Visualization**: Visual feedback for analysis results +- **Confidence Scoring**: Reliability indicators for all outputs + +## Installation + +1. Copy this folder to your ComfyUI `custom_nodes` directory: + ``` + ComfyUI/custom_nodes/NormalMapLightEstimator/ + ``` + +2. Restart ComfyUI + +3. Dependencies will install automatically. If installation fails, see troubleshooting. + +## Usage + +### Basic Setup + +1. Connect a normal map to the `normal_map` input +2. Connect a luma/lighting image to the `luma_image` input +3. Adjust the `luma_threshold` to control which areas are considered lit +4. Run the workflow to get directional and quality outputs + +### Inputs + +| Name | Type | Description | Default | +|------|------|-------------|---------| +| `normal_map` | IMAGE | RGB normal map in tangent space (0–255) | - | +| `luma_image` | IMAGE | Grayscale or RGB image to extract luminance | - | +| `luma_threshold` | FLOAT | Threshold to mask out non-lit areas | 0.5 | +| `curve_type` | COMBO | Luma processing curve | s_curve | +| `x_threshold` | FLOAT | Sensitivity for left/right detection | 0.1 | +| `y_threshold` | FLOAT | Sensitivity for top/bottom detection | 0.1 | +| `central_threshold` | FLOAT | Threshold for central lighting detection | 0.3 | +| `hard_light_threshold` | FLOAT | Below this = hard light | 0.15 | +| `soft_light_threshold` | FLOAT | Above this = soft light | 0.35 | +| `use_weighted` | BOOLEAN | Use weighted averaging | False | +| `custom_curve_points` | STRING | Custom curve points (optional) | - | + +### Outputs + +| Name | Type | Description | +|------|------|-------------| +| `x_direction` | STRING | "left", "central", "right" | +| `y_direction` | STRING | "top", "central", "bottom" | +| `combined_direction` | STRING | "top-left", "top-central", etc. | +| `hard_soft_index` | FLOAT | 0.0 (hard) to 1.0 (soft) | +| `x_confidence` | FLOAT | Confidence in X direction classification | +| `y_confidence` | FLOAT | Confidence in Y direction classification | +| `overall_confidence` | FLOAT | Combined confidence score | +| `spread_value` | FLOAT | Raw numerical spread value | +| `debug_mask` | IMAGE | Binary mask showing which normals were used | +| `directional_viz` | IMAGE | Color-coded visualization of results | + +## Algorithm + +### Core Concept +The system analyzes surface normals to infer lighting information: +1. **Normal Map (RGB)** → Surface orientation (XYZ) +2. **Luma Mask** → Which areas are actually lit +3. **Filtered Normals** → Only consider lit surface orientations +4. **Statistical Analysis** → Infer light direction and quality + +### Key Insight +Lit surface normals point **TOWARD** the light source. By averaging these directions, we get the light direction. + +### Processing Pipeline +1. **Normal Processing**: Convert RGB to XYZ, apply axis inversions +2. **Luma Masking**: Extract luminance, apply threshold with curves +3. **Directional Analysis**: Quadrant-based normal distribution analysis +4. **Quality Classification**: Spread analysis for hard/soft determination + +## Curve Types + +### Linear +Simple threshold-like behavior for basic masking. + +### S-curve +Enhance mid-tones, compress highlights/shadows. Good for balanced lighting. + +### Exponential +Emphasize bright areas. Useful for high-contrast scenes. + +### Logarithmic +Compress highlights, expand shadows. Good for low-light scenarios. + +### Custom +User-defined curve points for advanced control. + +## Threshold Controls + +### Directional Thresholds +- **X Threshold**: Controls left/right detection sensitivity +- **Y Threshold**: Controls top/bottom detection sensitivity +- **Central Threshold**: Detects central lighting patterns + +### Quality Thresholds +- **Hard Light Threshold**: Below this = hard light (index 0.0) +- **Soft Light Threshold**: Above this = soft light (index 1.0) +- **Between**: Linear interpolation from 0.0 to 1.0 + +## Usage Examples + +### Basic Directional Detection +``` +X Threshold: 0.1 (sensitive to left/right) +Y Threshold: 0.1 (sensitive to top/bottom) +Central Threshold: 0.3 (detect central lighting) +``` + +### Hard Light Detection +``` +Hard Threshold: 0.1 (very sensitive) +Soft Threshold: 0.2 (narrow range) +Result: Index 0.0-0.5 range +``` + +### Soft Light Detection +``` +Hard Threshold: 0.2 (less sensitive) +Soft Threshold: 0.5 (wider range) +Result: Index 0.5-1.0 range +``` + +## Visualization System + +### Debug Mask +- **White pixels**: Lit areas used in analysis +- **Black pixels**: Masked out areas + +### Directional Visualization +- **Red Channel**: X direction (left=255, central=192, right=128) +- **Green Channel**: Y direction (top=255, central=192, bottom=128) +- **Blue Channel**: Hard/soft index (hard=255, soft=128) +- **Confidence**: Applied as transparency effect + +## Troubleshooting + +### Common Issues + +**"No light direction detected"** +- Check if luma threshold is too high +- Verify normal map has proper lighting information +- Try adjusting directional thresholds + +**"Incorrect direction classification"** +- Try inverting X or Y axes for different normal conventions +- Adjust central threshold for omnidirectional lighting +- Check confidence scores for reliability + +**"Poor quality classification"** +- Adjust hard/soft thresholds based on your use case +- Try different curve types for luma processing +- Use weighted averaging for more accurate results + +### Manual Installation + +If auto-installation fails: + +```bash +cd ComfyUI/custom_nodes/NormalMapLightEstimator/ +pip install -r requirements.txt +``` + +## Technical Details + +### Dependencies +- torch>=1.9.0 +- numpy>=1.21.0 +- Pillow>=8.0.0 +- scipy>=1.7.0 + +### Performance +- Optimized for ComfyUI tensor operations +- Batch processing support +- Memory-efficient covariance calculations +- GPU acceleration where possible + +### Edge Cases Handled +- No lit areas: Returns default central direction +- Single normal: Uses that direction +- Collinear normals: Handles degenerate cases gracefully +- Low confidence: Provides reliability indicators + +## License + +MIT License - see LICENSE for details. + +## Changelog + +### v1.0.0 +- Initial release with categorical directional outputs +- Hard/soft classification with continuous index +- Curves strategy for luma processing +- User-controllable thresholds +- Debug visualization system +- Confidence scoring system + +--- + +For support, open an issue on GitHub. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..13e1160 --- /dev/null +++ b/__init__.py @@ -0,0 +1,4 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +# Re-export the mappings for ComfyUI +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/example/README.md b/example/README.md new file mode 100644 index 0000000..d6a57a3 --- /dev/null +++ b/example/README.md @@ -0,0 +1,39 @@ +# Example Workflows + +This directory contains example workflows demonstrating the NormalMapLightEstimator node. + +## Files + +- `example_workflow.json` - Basic workflow showing node usage +- `README.md` - This file + +## Usage + +1. Load the example workflow in ComfyUI +2. Connect your normal map and luma image +3. Adjust parameters as needed +4. Run the workflow to see results + +## Example Parameters + +### Basic Setup +- `luma_threshold`: 0.5 (default) +- `curve_type`: "s_curve" (enhances mid-tones) +- `x_threshold`: 0.1 (sensitive to left/right) +- `y_threshold`: 0.1 (sensitive to top/bottom) + +### Hard Light Detection +- `hard_light_threshold`: 0.1 (very sensitive) +- `soft_light_threshold`: 0.2 (narrow range) + +### Soft Light Detection +- `hard_light_threshold`: 0.2 (less sensitive) +- `soft_light_threshold`: 0.5 (wider range) + +## Outputs + +The node provides: +- Categorical directions (X, Y, combined) +- Hard/soft index (0.0 to 1.0) +- Confidence scores +- Debug visualizations diff --git a/example/example_workflow.json b/example/example_workflow.json new file mode 100644 index 0000000..0032198 --- /dev/null +++ b/example/example_workflow.json @@ -0,0 +1,31 @@ +{ + "workflow": { + "1": { + "class_type": "LoadImage", + "inputs": { + "image": "normal_map.png" + } + }, + "2": { + "class_type": "LoadImage", + "inputs": { + "image": "luma_image.png" + } + }, + "3": { + "class_type": "NormalMapLightEstimator", + "inputs": { + "normal_map": ["1", 0], + "luma_image": ["2", 0], + "luma_threshold": 0.5, + "curve_type": "s_curve", + "x_threshold": 0.1, + "y_threshold": 0.1, + "central_threshold": 0.3, + "hard_light_threshold": 0.15, + "soft_light_threshold": 0.35, + "use_weighted": false + } + } + } +} diff --git a/nodes/__init__.py b/nodes/__init__.py new file mode 100644 index 0000000..cdd2cf9 --- /dev/null +++ b/nodes/__init__.py @@ -0,0 +1,11 @@ +from .nodes import NormalMapLightEstimator + +NODE_CLASS_MAPPINGS = { + "NormalMapLightEstimator": NormalMapLightEstimator +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "NormalMapLightEstimator": "Normal Map Light Estimator" +} + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/nodes/nodes.py b/nodes/nodes.py new file mode 100644 index 0000000..0e5e12f --- /dev/null +++ b/nodes/nodes.py @@ -0,0 +1,189 @@ +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 + +class NormalMapLightEstimator: + """ + ComfyUI custom node for estimating light direction and quality from normal maps. + """ + + CATEGORY = "Custom/Lighting" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "normal_map": ("IMAGE",), + "luma_image": ("IMAGE",), + "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}), + }, + "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 + ) + + RETURN_NAMES = ( + "x_direction", + "y_direction", + "combined_direction", + "hard_soft_index", + "x_confidence", + "y_confidence", + "overall_confidence", + "spread_value", + "debug_mask", + "directional_viz" + ) + + 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=""): + """ + Enhanced categorical light direction estimation with hard/soft index. + """ + # 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 + ) + + # Extract results for each batch + x_directions = [] + y_directions = [] + combined_directions = [] + hard_soft_indices = [] + x_confidences = [] + y_confidences = [] + overall_confidences = [] + spread_values = [] + + 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']) + + # Generate debug outputs + debug_mask = DebugVisualizer.generate_debug_mask(mask) + directional_viz = DebugVisualizer.generate_directional_visualization( + normal_map, categorical_results + ) + + 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 + ) + + 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/requirements.txt b/requirements.txt new file mode 100644 index 0000000..9e6784e --- /dev/null +++ b/requirements.txt @@ -0,0 +1,4 @@ +torch>=1.9.0 +numpy>=1.21.0 +Pillow>=8.0.0 +scipy>=1.7.0 diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..750411e --- /dev/null +++ b/setup.py @@ -0,0 +1,32 @@ +from setuptools import setup, find_packages + +with open("README.md", "r", encoding="utf-8") as fh: + long_description = fh.read() + +with open("requirements.txt", "r", encoding="utf-8") as fh: + requirements = [line.strip() for line in fh if line.strip() and not line.startswith("#")] + +setup( + name="comfyui-light-direction-estimation", + version="1.0.0", + author="ComfyUI Light Direction Estimation", + description="A ComfyUI custom node for estimating light direction and quality from normal maps", + long_description=long_description, + long_description_content_type="text/markdown", + url="https://github.com/your-username/Comfyui-LightDirection-estimation", + packages=find_packages(), + classifiers=[ + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "License :: OSI Approved :: MIT License", + "Operating System :: OS Independent", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + ], + python_requires=">=3.8", + install_requires=requirements, + include_package_data=True, +) diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 0000000..947b7dc --- /dev/null +++ b/utils/__init__.py @@ -0,0 +1,5 @@ +from .light_estimator import CategoricalLightEstimator +from .luma_mask_processor import LumaMaskProcessor +from .debug_visualizer import DebugVisualizer + +__all__ = ['CategoricalLightEstimator', 'LumaMaskProcessor', 'DebugVisualizer'] diff --git a/utils/debug_visualizer.py b/utils/debug_visualizer.py new file mode 100644 index 0000000..3bf2e24 --- /dev/null +++ b/utils/debug_visualizer.py @@ -0,0 +1,236 @@ +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) diff --git a/utils/light_estimator.py b/utils/light_estimator.py new file mode 100644 index 0000000..3447e96 --- /dev/null +++ b/utils/light_estimator.py @@ -0,0 +1,268 @@ +import torch +import numpy as np + +class CategoricalLightEstimator: + """ + Light estimator with categorical directional outputs and simplified hard/soft classification. + """ + + def __init__(self): + # Directional thresholds + self.x_threshold = 0.1 + self.y_threshold = 0.1 + self.central_threshold = 0.3 + + # Hard/Soft classification thresholds + self.hard_light_threshold = 0.15 # Below this = hard light + self.soft_light_threshold = 0.35 # Above this = soft light + + def analyze_directional_categories(self, normals, mask): + """ + Analyze normals and classify into directional categories with hard/soft index. + + Args: + normals: Normal vectors tensor (B, H, W, 3) + mask: Binary mask tensor (B, H, W) + + Returns: + results: List of analysis results for each batch + """ + batch_size = normals.shape[0] + results = [] + + for b in range(batch_size): + lit_normals = normals[b][mask[b]] + + if len(lit_normals) == 0: + results.append(self._empty_categories()) + continue + + # 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 + }) + + return results + + 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 + } + + # 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 confidence based on how clear the spread is + confidence = min(1.0, spread / 0.5) # Normalize confidence + + return { + 'spread': spread, + 'confidence': confidence + } + + 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 + + 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 + """ + x_components = xy_normals[:, 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() + + 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 + } + + 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 + """ + y_components = xy_normals[:, 1] + + # 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 + } + + 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" + """ + 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): + return "central" + elif left_pct > 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" + """ + 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): + return "central" + elif top_pct > bottom_pct: + return "top" + else: + return "bottom" + + def _empty_categories(self): + """ + 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}, + 'quality_analysis': {'spread': 0.0, 'confidence': 0.0} + } diff --git a/utils/luma_mask_processor.py b/utils/luma_mask_processor.py new file mode 100644 index 0000000..28c8f47 --- /dev/null +++ b/utils/luma_mask_processor.py @@ -0,0 +1,160 @@ +import torch +import numpy as np +from scipy.interpolate import interp1d + +class LumaMaskProcessor: + """ + Enhanced luma mask processor with curves strategy and multiple processing modes. + """ + + @staticmethod + def process_with_curves(luma_image, curve_type="linear", curve_points=None, threshold=0.5): + """ + Process luma with various curve strategies. + + Args: + luma_image: Input luma values tensor (B, H, W, C) + curve_type: "linear", "s_curve", "exponential", "logarithmic", "custom" + curve_points: Custom curve points for "custom" type + threshold: Final threshold after curve processing + + Returns: + mask: Binary mask tensor (B, H, W) where True indicates lit areas + """ + # Convert to numpy for curve processing + luma_np = luma_image.detach().cpu().numpy() + + if curve_type == "linear": + # Simple linear mapping + processed_luma = luma_np + + elif curve_type == "s_curve": + # S-curve: enhance mid-tones + processed_luma = 3 * luma_np**2 - 2 * luma_np**3 + + elif curve_type == "exponential": + # Exponential: emphasize bright areas + processed_luma = luma_np**2 + + elif curve_type == "logarithmic": + # Logarithmic: compress highlights + processed_luma = np.log(1 + luma_np * 9) / np.log(10) + + elif curve_type == "custom" and curve_points: + # Custom curve interpolation + x_points, y_points = zip(*curve_points) + curve_func = interp1d(x_points, y_points, kind='cubic', + bounds_error=False, fill_value='extrapolate') + processed_luma = curve_func(luma_np) + + else: + processed_luma = luma_np + + # Apply threshold + mask = processed_luma > threshold + + return torch.from_numpy(mask).to(luma_image.device) + + @staticmethod + def weighted_luma_mask(luma_image, threshold=0.5): + """ + Create weighted mask using luma intensity as weights. + + Args: + luma_image: RGB or grayscale image tensor (B, H, W, C) + threshold: Minimum luma threshold + + Returns: + weights: Weight tensor (B, H, W) with luma values as weights + """ + if luma_image.shape[-1] == 3: # RGB image + luma = 0.2126 * luma_image[..., 0] + 0.7152 * luma_image[..., 1] + 0.0722 * luma_image[..., 2] + else: # Already grayscale + luma = luma_image[..., 0] if luma_image.shape[-1] == 1 else luma_image + + # Use luma values as weights, zero out below threshold + weights = torch.where(luma > threshold, luma, torch.zeros_like(luma)) + + return weights + + @staticmethod + def adaptive_threshold(luma_image, method="otsu"): + """ + Apply adaptive thresholding methods. + + Args: + luma_image: Input luma values tensor (B, H, W, C) + method: "otsu", "mean", "median" + + Returns: + mask: Binary mask tensor (B, H, W) + """ + luma_np = luma_image.detach().cpu().numpy() + + if method == "otsu": + # Otsu's method for optimal threshold + try: + from skimage.filters import threshold_otsu + threshold = threshold_otsu(luma_np) + except ImportError: + # Fallback to mean if skimage not available + threshold = np.mean(luma_np) + + elif method == "mean": + # Mean-based threshold + threshold = np.mean(luma_np) + + elif method == "median": + # Median-based threshold + threshold = np.median(luma_np) + + else: + threshold = 0.5 + + mask = luma_np > threshold + return torch.from_numpy(mask).to(luma_image.device) + + @staticmethod + def multi_scale_processing(luma_image, scales=[1.0, 0.5, 2.0], weights=[0.5, 0.3, 0.2]): + """ + Process luma at multiple scales and combine results. + + Args: + luma_image: Input luma values tensor (B, H, W, C) + scales: List of scale factors + weights: Weights for combining scales + + Returns: + combined_mask: Combined binary mask tensor (B, H, W) + """ + results = [] + + for scale, weight in zip(scales, weights): + # Resize image + if scale != 1.0: + resized = torch.nn.functional.interpolate( + luma_image.permute(0, 3, 1, 2), + scale_factor=scale, + mode='bilinear', + align_corners=False + ).permute(0, 2, 3, 1) + else: + resized = luma_image + + # Process at this scale + processed = LumaMaskProcessor.process_with_curves(resized, "s_curve") + + # Resize back to original size + if scale != 1.0: + processed = torch.nn.functional.interpolate( + processed.unsqueeze(1).float(), + size=luma_image.shape[1:3], + mode='bilinear', + align_corners=False + ).squeeze(1) + + results.append(processed * weight) + + # Combine results + combined = torch.sum(torch.stack(results), dim=0) + return combined > 0.5