add: PercentileToPixelValueTonemap

PercentileToPixelValueTonemap.py

@@ -0,0 +1,85 @@
+import numpy as np
+import torch
+
+class PercentileToPixelValueTonemap:
+    """
+    DiffusionLight Percentile map to pixel value tonemap 
+    """
+    def __init__(self):
+        pass 
+    
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+                "percentile": ("FLOAT", {
+                    "default": 90,
+                    "min": 0,
+                    "max": 100,
+                    "step": 0.01,
+                    "round": False,
+                    "display": "number",
+                    "lazy": True
+                }),
+                "pixel_value": ("FLOAT", {
+                    "default": 0.9,
+                    "min": 0,
+                    "max": 1,
+                    "step": 0.000001,
+                    "round": False,
+                    "display": "number",
+                    "lazy": True
+                }),
+                "gamma": ("FLOAT", {
+                    "default": 2.4,
+                    "min": -1000,
+                    "max": 1000,
+                    "step": 0.01,
+                    "round": False,
+                    "display": "number",
+                    "lazy": True
+                }),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+
+    FUNCTION = "percentile_to_pixel_value_tonemap"
+
+    def percentile_to_pixel_value_tonemap(self, images, percentile, pixel_value, gamma):
+        """
+        map percentile of the HDR image to some value for tonemapping
+        set gamma to 1.0 to disable gamma correction
+        Args:
+            images (IMAGE): The input HDR image. #Tensor of image format shape (range 0-1) shape [N, H, W, 3]
+            percentile (float): The percentile value to use for tonemapping.
+            pixel_value (float): The pixel value to map the percentile to.
+            gamma (float): The gamma value to apply during the conversion.
+        """
+        # apply gamma correction
+        if gamma != 1.0:
+            images = torch.pow(images, 1.0 / gamma)
+
+        # calculate the percentile value in beach image in batch dimension
+        percentile_value = batch_percentile(images, percentile)
+
+        # map the percentile value to the pixel value
+        hdr_image = images / percentile_value[:,None,None,None] * pixel_value
+
+        return (hdr_image, )
+
+def batch_percentile(input_tensor: torch.Tensor, percentile: float) -> torch.Tensor:
+    """
+    input_tensor: shape [b, H, W, 3]
+    percentile: scalar float between 0 and 100
+    returns: shape [b]
+    """
+    # Flatten [H, W, 3] -> [-1] per image
+    flattened = input_tensor.view(input_tensor.shape[0], -1)  # shape [b, H*W*3]
+
+    # Compute percentile along dimension 1 (per image)
+    result = torch.quantile(flattened, percentile / 100.0, dim=1)
+    return result
+    
+

__init__.py

@@ -4,6 +4,7 @@ from .SaveHDR import SaveHDR
 from .ExposureBracket import ExposureBracket
 from .PadBlackBorder import PadBlackBorder
 from .ChromeballMask import ChromeballMask
+from .PercentileToPixelValueTonemap import PercentileToPixelValueTonemap
 
 
 # A dictionary that contains all nodes you want to export with their names
@@ -15,6 +16,7 @@ NODE_CLASS_MAPPINGS = {
     "DiffusionLightExposureBracket": ExposureBracket,
     "DiffusionLightPadBlackBorder": PadBlackBorder,
     "DiffusionLightChromeballMask": ChromeballMask,
+    "DiffusionLightPercentileToPixelValueTonemap": PercentileToPixelValueTonemap,
 }
 
 # A dictionary that contains the friendly/humanly readable titles for the nodes
@@ -25,6 +27,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "DiffusionLightExposureBracket": "ExposureBracket",
     "DiffusionLightPadBlackBorder": "PadBlackBorder",
     "DiffusionLightChromeballMask": "ChromeballMask",
+    "DiffusionLightPercentileToPixelValueTonemap": "PercentileToPixelValueTonemap",
 }