From 9c18b55439bac8067aa37a2c8aef80c95d27daed Mon Sep 17 00:00:00 2001
From: xander <xandersteenbruggex@gmail.com>
Date: Thu, 26 Sep 2024 21:14:19 +0200
Subject: [PATCH] add depth parallax mode for the z-direction

---
 nodes/images/flex_images.py | 38 ++++++++++++++++++++++++++++++++-----
 1 file changed, 33 insertions(+), 5 deletions(-)

diff --git a/nodes/images/flex_images.py b/nodes/images/flex_images.py
index a1cbeda..96b69fb 100644
--- a/nodes/images/flex_images.py
+++ b/nodes/images/flex_images.py
@@ -377,27 +377,55 @@ class FlexImageParallax(FlexImageBase):
                 **super().INPUT_TYPES()["required"],
                 "shift_x": ("FLOAT", {"default": 0.1, "min": -1.0, "max": 1.0, "step": 0.01}),
                 "shift_y": ("FLOAT", {"default": 0.1, "min": -1.0, "max": 1.0, "step": 0.01}),
+                "shift_z": ("FLOAT", {"default": 0.1, "min": -1.0, "max": 1.0, "step": 0.01}),
                 "depth_map": ("IMAGE",),  # Assuming depth map is provided as an image
             }
         }
 
     @classmethod
     def get_modifiable_params(cls):
-        return ["shift_x", "shift_y", "None"]
+        return ["shift_x", "shift_y", "shift_z"]
 
-    def apply_effect_internal(self, image: np.ndarray, shift_x: float, shift_y: float, depth_map: np.ndarray, frame_index: int, **kwargs) -> np.ndarray:
+    def apply_effect_internal(
+        self,
+        image: np.ndarray,
+        shift_x: float,
+        shift_y: float,
+        shift_z: float,
+        depth_map: np.ndarray,
+        frame_index: int,
+        **kwargs
+    ) -> np.ndarray:
         h, w, _ = image.shape
 
+        # Extract the corresponding depth map frame and normalize it
         depth_map_frame = depth_map[frame_index].cpu().numpy()
         depth_map_gray = np.mean(depth_map_frame, axis=-1)
-        depth_map_gray = depth_map_gray / np.max(depth_map_gray)
+        depth_map_gray /= np.max(depth_map_gray)
+
+        # Calculate shifts in x and y directions
         dx = (w * shift_x * depth_map_gray).astype(np.int32)
         dy = (h * shift_y * depth_map_gray).astype(np.int32)
+
+        # Calculate scaling factor for z-direction parallax
+        scale = 1 + shift_z * depth_map_gray
+
         x, y = np.meshgrid(np.arange(w), np.arange(h))
 
-        new_x = np.clip(x + dx, 0, w - 1)
-        new_y = np.clip(y + dy, 0, h - 1)
+        # Apply shifts
+        x_shifted = x + dx
+        y_shifted = y + dy
 
+        # Apply scaling around the center of the image
+        cx, cy = w / 2, h / 2
+        x_scaled = cx + (x_shifted - cx) * scale
+        y_scaled = cy + (y_shifted - cy) * scale
+
+        # Ensure coordinates are within image bounds
+        new_x = np.clip(x_scaled, 0, w - 1).astype(np.int32)
+        new_y = np.clip(y_scaled, 0, h - 1).astype(np.int32)
+
+        # Generate the resulting image with parallax effect
         result = image[new_y, new_x]
 
         return np.clip(result, 0, 1)
\ No newline at end of file