Create displacement_map_node.py

displacement_map_node.py

@@ -0,0 +1,55 @@
+import torch
+import numpy as np
+from PIL import Image
+import torchvision.transforms as transforms
+
+class ExtractDisplacementMap:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE",),  # Input image (e.g., normal or bump map)
+                "intensity": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1}),  # Displacement intensity
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("displacement_map",)
+    FUNCTION = "extract_displacement_map"
+    CATEGORY = "image/processing"
+    OUTPUT_NODE = False
+
+    def extract_displacement_map(self, image, intensity):
+        # Convert ComfyUI image tensor to numpy array
+        if isinstance(image, torch.Tensor):
+            image = image.cpu().numpy()
+        if image.ndim == 4:
+            image = image[0]  # Remove batch dimension if present
+        if image.shape[0] in [1, 3]:
+            image = np.transpose(image, (1, 2, 0))  # Convert (C, H, W) to (H, W, C)
+        image = (image * 255).astype(np.uint8)  # Scale to 0-255
+
+        # Convert to grayscale
+        if image.shape[-1] == 3:
+            image = np.dot(image[..., :3], [0.2989, 0.5870, 0.1140])  # RGB to grayscale
+        else:
+            image = image[..., 0]  # Use the first channel if already grayscale
+
+        # Normalize and scale by intensity
+        displacement_map = image.astype(np.float32) / 255.0
+        displacement_map = displacement_map * intensity
+
+        # Convert to tensor and add batch dimension
+        displacement_map = torch.from_numpy(displacement_map).unsqueeze(0).unsqueeze(0)
+        displacement_map = displacement_map.to(torch.float32)
+
+        return (displacement_map,)
+
+# Node mappings
+NODE_CLASS_MAPPINGS = {
+    "ExtractDisplacementMap": ExtractDisplacementMap
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "ExtractDisplacementMap": "Extract Displacement Map"
+}