Add files via upload

2024-04-05 17:51:27 +09:00
parent e4fed55426
commit 5675f03330
1 changed files with 109 additions and 0 deletions
@@ -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"
+}
+
+
+