diff --git a/PercentileToPixelValueTonemap.py b/PercentileToPixelValueTonemap.py new file mode 100644 index 0000000..eddc02b --- /dev/null +++ b/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 + + diff --git a/__init__.py b/__init__.py index bfa83cc..c327ea4 100644 --- a/__init__.py +++ b/__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", }