From 5675f0333091514af491d78de2ece9583fc38642 Mon Sep 17 00:00:00 2001 From: TJ16th Date: Fri, 5 Apr 2024 17:51:27 +0900 Subject: [PATCH] Add files via upload --- EulerLightingNode.py | 109 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 109 insertions(+) create mode 100644 EulerLightingNode.py diff --git a/EulerLightingNode.py b/EulerLightingNode.py new file mode 100644 index 0000000..c186629 --- /dev/null +++ b/EulerLightingNode.py @@ -0,0 +1,109 @@ + + +import numpy as np +from PIL import Image +import torch + +class EulerLightingNode: + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "diffuse_map": ("IMAGE",), + "normal_map": ("IMAGE",), + "specular_map": ("IMAGE",), + "light_yaw": ("FLOAT", {"default": 45, "min": -180, "max": 180, "step": 1}), + "light_pitch": ("FLOAT", {"default": 30, "min": -90, "max": 90, "step": 1}), + "specular_power": ("FLOAT", {"default": 32, "min": 1, "max": 200, "step": 1}), + "ambient_light": ("FLOAT", {"default": 0.50, "min": 0, "max": 1, "step": 0.01}), + "NormalDiffuseStrength": ("FLOAT", {"default": 1.00, "min": 0, "max": 5.0, "step": 0.01}), + "SpecularHighlightsStrength": ("FLOAT", {"default": 1.00, "min": 0, "max": 5.0, "step": 0.01}), + "TotalGain": ("FLOAT", {"default": 1.00, "min": 0, "max": 2.0, "step": 0.01}), + }, + } + + RETURN_TYPES = ("IMAGE",) + + FUNCTION = "execute" + CATEGORY = "custom" + + def execute(self, diffuse_map, normal_map, specular_map, light_yaw, light_pitch, specular_power, ambient_light,NormalDiffuseStrength,SpecularHighlightsStrength,TotalGain): + + # 入力画像をテンソルに変換 + diffuse_tensor = diffuse_map.permute(0, 3, 1, 2) # (1, 512, 512, 3) -> (1, 3, 512, 512) + normal_tensor = normal_map.permute(0, 3, 1, 2) * 2.0 - 1.0 # (1, 512, 512, 3) -> (1, 3, 512, 512) + specular_tensor = specular_map.permute(0, 3, 1, 2) # (1, 512, 512, 3) -> (1, 3, 512, 512) + + # 法線ベクトルを正規化 + normal_tensor = torch.nn.functional.normalize(normal_tensor, dim=1) + + # light_directionをブロードキャスト用に正しくリシェイプ + light_direction = self.euler_to_vector(light_yaw, light_pitch, 0 ) + light_direction = light_direction.view(1, 3, 1, 1) # [1, 3, 1, 1]にリシェイプしてブロードキャストを可能にする + + # camera_directionをブロードキャスト用に正しくリシェイプ + camera_direction = self.euler_to_vector(0,0,0) + camera_direction = camera_direction.view(1, 3, 1, 1) # [1, 3, 1, 1]にリシェイプしてブロードキャストを可能にする + + # 乗算のための既存のコード... + diffuse = torch.sum(normal_tensor * light_direction, dim=1, keepdim=True) + diffuse = torch.clamp(diffuse, 0, 1) + + # 鏡面反射の計算 + half_vector = torch.nn.functional.normalize(light_direction + camera_direction, dim=1) + specular = torch.sum(normal_tensor * half_vector, dim=1, keepdim=True) + specular = torch.pow(torch.clamp(specular, 0, 1), specular_power) + + # 拡散反射と鏡面反射の結果を合成 + output_tensor = ( diffuse_tensor * (ambient_light + diffuse * NormalDiffuseStrength ) + specular_tensor * specular * SpecularHighlightsStrength) * TotalGain + + # テンソルを出力用の画像に変換 + output_tensor = output_tensor.permute(0, 2, 3, 1) # (1, 3, 512, 512) -> (1, 512, 512, 3) + + return (output_tensor,) + + + def euler_to_vector(self, yaw, pitch, roll): + yaw_rad = np.radians(yaw) + pitch_rad = np.radians(pitch) + roll_rad = np.radians(roll) + + cos_pitch = np.cos(pitch_rad) + sin_pitch = np.sin(pitch_rad) + cos_yaw = np.cos(yaw_rad) + sin_yaw = np.sin(yaw_rad) + cos_roll = np.cos(roll_rad) + sin_roll = np.sin(roll_rad) + + direction = np.array([ + sin_yaw * cos_pitch, + sin_pitch, + cos_pitch * cos_yaw + ]) + + + return torch.from_numpy(direction).float() + + def convert_tensor_to_image(self, tensor): + # PyTorchのテンソルをPIL.Imageに変換 + tensor = tensor.squeeze(0) # (1, 512, 512, 3) -> (512, 512, 3) + tensor = tensor.clamp(0, 1) # テンソルの値を0から1の範囲に制限 + image = Image.fromarray((tensor.detach().cpu().numpy() * 255).astype(np.uint8)) + return image + + + +NODE_CLASS_MAPPINGS = { + "EulerLightingNode": EulerLightingNode +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "EulerLightingNode": "NormalLighting" +} + + +