Initial commit

2024-04-02 19:38:51 +03:00
parent fbdb933e6f
commit 8f0d05f67c
5 changed files with 3853 additions and 0 deletions
@@ -0,0 +1,160 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
@@ -0,0 +1,3 @@
+from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
@@ -0,0 +1,153 @@
+import torch
+import numpy as np
+import torch
+import torch.nn.functional as F
+from .relighting.tonemapper import TonemapHDR
+
+def create_envmap_grid(size: int):
+    """
+    BLENDER CONVENSION
+    Create the grid of environment map that contain the position in sperical coordinate
+    Top left is (0,0) and bottom right is (pi/2, 2pi)
+    """    
+    theta = torch.linspace(0, np.pi * 2, size * 2)
+    phi = torch.linspace(0, np.pi, size)
+    
+    #use indexing 'xy' torch match vision's homework 3
+    theta, phi = torch.meshgrid(theta, phi ,indexing='xy') 
+    
+    theta_phi = torch.cat([theta[..., None], phi[..., None]], dim=-1)
+    theta_phi = theta_phi.numpy()
+    return theta_phi
+
+def get_normal_vector(incoming_vector: np.ndarray, reflect_vector: np.ndarray):
+    """
+    BLENDER CONVENSION
+    incoming_vector: the vector from the point to the camera
+    reflect_vector: the vector from the point to the light source
+    """
+    #N = 2(R ⋅ I)R - I
+    N = (incoming_vector + reflect_vector) / np.linalg.norm(incoming_vector + reflect_vector, axis=-1, keepdims=True)
+    return N
+
+def get_cartesian_from_spherical(theta: np.array, phi: np.array, r = 1.0):
+    """
+    BLENDER CONVENSION
+    theta: vertical angle
+    phi: horizontal angle
+    r: radius
+    """
+    x = r * np.sin(theta) * np.cos(phi)
+    y = r * np.sin(theta) * np.sin(phi)
+    z = r * np.cos(theta)
+    return np.concatenate([x[...,None],y[...,None],z[...,None]], axis=-1)
+
+class chrome_ball_to_envmap:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "ball_images": ("IMAGE", ),
+                "envmap_height": ("INT", {"default": 256, "min": 1, "max": 2048, "step": 1}, ),
+                "scale": ("INT", {"default": 4, "min": 1, "max": 30, "step": 1}, ),
+            },
+        }
+        
+    CATEGORY = "DiffusionLight"
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("image",  )
+
+    FUNCTION = "process"
+
+   
+    def process(self, ball_images, envmap_height, scale):     
+        I = np.array([1, 0, 0])
+
+        # compute  normal map that create from reflect vector
+        env_grid = create_envmap_grid(envmap_height * scale)   
+        reflect_vec = get_cartesian_from_spherical(env_grid[...,1], env_grid[...,0])
+        normal = get_normal_vector(I[None,None], reflect_vec)
+        
+        # turn from normal map to position to lookup [Range: 0,1]
+        pos = (normal + 1.0) / 2
+        pos  = 1.0 - pos
+        pos = pos[...,1:]
+        
+        env_map = None
+        # convert position to pytorch grid look up
+        grid = torch.from_numpy(pos)[None].float()
+        grid = grid * 2 - 1 # convert to range [-1,1]
+        print(grid.shape)
+        ball_images = ball_images.permute(0,3,1,2) # [1,3,H,W]
+
+        env_maps_list = []
+        for ball in ball_images:
+            env_map = F.grid_sample(ball.unsqueeze(0), grid, mode='bilinear', padding_mode='border', align_corners=True)
+            env_map_default = F.interpolate(env_map, size=(envmap_height, envmap_height*2), mode='bilinear', align_corners=True)
+            env_map_default = env_map_default.permute(0,2,3,1).cpu().to(torch.float32)
+            env_maps_list.append(env_map_default)
+        env_maps_out = torch.cat(env_maps_list, dim=0)
+
+        
+ 
+        return env_maps_out,
+
+class exposure_to_hdr:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "images": ("IMAGE", ),
+                #"EV": ("FLOAT", {"default": 0, "min": 1, "max": 30, "step": 1}, ),
+                "gamma": ("FLOAT", {"default": 2.4, "min": 1, "max": 30, "step": 0.01}, ),
+            },
+        }
+        
+    CATEGORY = "DiffusionLight"
+    RETURN_TYPES = ("IMAGE", "IMAGE",)
+    RETURN_NAMES = ("hrd_image", "ldr_image", )
+
+    FUNCTION = "exposuretohdr"
+
+    def exposuretohdr(self, images, gamma):     
+        first_image = torch.pow(images[0], gamma)
+        evs = [0.0, -2.5, -5.0]
+        hdr2ldr = TonemapHDR(gamma=gamma, percentile=99, max_mapping=0.9)
+        scaler = torch.tensor([0.212671, 0.715160, 0.072169])
+        
+        # read luminace for every image
+        luminances = []
+        for i in range(len(evs)):
+            linear_img = torch.pow(images[i], gamma)
+            linear_img = linear_img * 1 / (2** evs[i])
+            # compute luminace
+            lumi = linear_img @ scaler
+            luminances.append(lumi)
+
+        # start from darkest image
+        out_luminace = luminances[len(evs) - 1]
+        for i in range(len(evs) - 1, 0, -1):
+            # compute mask
+            maxval = 1 / (2 ** evs[i-1])
+            p1 = torch.clip((luminances[i-1] - 0.9 * maxval) / (0.1 * maxval), 0, 1)
+            p2 = out_luminace > luminances[i-1]
+            mask = (p1 * p2)
+            out_luminace = luminances[i-1] * (1-mask) + out_luminace * mask
+        
+        hdr_rgb = first_image * (out_luminace / (luminances[0] + 1e-10)).unsqueeze(-1)
+        ldr_rgb, _, _ = hdr2ldr(hdr_rgb)
+
+        hrd_rgb = hdr_rgb.unsqueeze(0).cpu().to(torch.float32)
+        ldr_rgb = ldr_rgb.unsqueeze(0).cpu().to(torch.float32)
+
+        return (hrd_rgb, ldr_rgb,)
+
+NODE_CLASS_MAPPINGS = {
+    "chrome_ball_to_envmap": chrome_ball_to_envmap,
+    "exposure_to_hdr": exposure_to_hdr,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "chrome_ball_to_envmap": "Chrome Ball to Envmap",    
+    "exposure_to_hdr": "Exposure to HDR",
+}
@@ -0,0 +1,37 @@
+import numpy as np
+
+import torch
+
+class TonemapHDR(object):
+    """
+        Tonemap HDR image globally. First, we find alpha that maps the (max(tensor_img) * percentile) to max_mapping.
+        Then, we calculate I_out = alpha * I_in ^ (1/gamma)
+        input : torch.Tensor batch of images : [H, W, C]
+        output : torch.Tensor batch of images : [H, W, C]
+    """
+
+    def __init__(self, gamma=2.4, percentile=50, max_mapping=0.5):
+        self.gamma = gamma
+        self.percentile = percentile
+        self.max_mapping = max_mapping  # the value to which alpha will map the (max(tensor_img) * percentile) to
+
+    def __call__(self, tensor_img, clip=True, alpha=None, gamma=True):
+        if gamma:
+            power_tensor_img = torch.pow(tensor_img, 1 / self.gamma)
+        else:
+            power_tensor_img = tensor_img
+        non_zero = power_tensor_img > 0
+        if non_zero.any():
+            r_percentile = torch.quantile(power_tensor_img[non_zero], self.percentile / 100.0)
+        else:
+            r_percentile = torch.quantile(power_tensor_img, self.percentile / 100.0)
+        if alpha is None:
+            alpha = self.max_mapping / (r_percentile + 1e-10)
+        tonemapped_img = alpha * power_tensor_img
+
+        if clip:
+            tonemapped_img_clip = torch.clamp(tonemapped_img, 0, 1)
+        else:
+            tonemapped_img_clip = tonemapped_img
+
+        return tonemapped_img_clip.float(), alpha, tonemapped_img