From d0130eecacbe4e828f98bbd0b1f9c3e50234ed5d Mon Sep 17 00:00:00 2001
From: provos <provos@gmail.com>
Date: Sun, 23 Jun 2024 13:04:38 -0700
Subject: [PATCH] refactor: Add is_ground_plane property to ImageSlice class

This is not currently working and maybe never work.
---
 slice.py | 64 +++++++++++++++++++++++++++++++++++++++++++-------------
 1 file changed, 50 insertions(+), 14 deletions(-)

diff --git a/slice.py b/slice.py
index e54953d..939b48d 100644
--- a/slice.py
+++ b/slice.py
@@ -8,17 +8,28 @@
 
 
 class ImageSlice:
-    __slots__ = ('image', '_depth', '_filename',
+    __slots__ = ('image', '_depth', '_filename', '_is_ground_plane',
                  'positive_prompt', 'negative_prompt')
 
     def __init__(self, image=None, depth=-1, filename=None, positive_prompt='', negative_prompt=''):
         self.image = image
         self._depth = depth
         self._filename = filename
+        self._is_ground_plane = False
 
         self.positive_prompt = positive_prompt
         self.negative_prompt = negative_prompt
 
+    @property
+    def is_ground_plane(self):
+        return self._is_ground_plane
+
+    @is_ground_plane.setter
+    def is_ground_plane(self, value):
+        if not isinstance(value, bool):
+            raise ValueError("is_ground_plane must be a boolean value")
+        self._is_ground_plane = value
+
     @property
     def depth(self):
         return self._depth
@@ -46,22 +57,47 @@ def __eq__(self, other):
         return (self.image == other.image and
                 self.depth == other.depth and
                 self.filename == other.filename)
-
-    def create_card(self, image_height: int, image_width: int, cam: Camera):
-        z = cam.max_distance * ((255 - self.depth) / 255.0)
-
+        
+    def _dimension_at_depth(self, z: float, image_height: int, image_width: int, cam: Camera):
         fl_px = cam.focal_length_px(image_width)
-
-        # Calculate the 3D points of the card corners
         card_width = (image_width * (z + cam.camera_distance)) / fl_px
         card_height = (image_height * (z + cam.camera_distance)) / fl_px
-
-        card_corners_3d = np.array([
-            [-card_width / 2, -card_height / 2, z],
-            [card_width / 2, -card_height / 2, z],
-            [card_width / 2, card_height / 2, z],
-            [-card_width / 2, card_height / 2, z]
-        ], dtype=np.float32)
+        return card_width, card_height
+    
+    def _depth_to_z(self, depth: float, cam: Camera):
+        return cam.max_distance * ((255 - depth) / 255.0)
+    
+    def create_card(self, image_height: int, image_width: int, cam: Camera):
+        if self.is_ground_plane:
+            near_z = 0
+            far_z = cam.max_distance
+            
+            near_width, near_height = self._dimension_at_depth(
+                near_z, image_height, image_width, cam)
+            far_width, far_height = self._dimension_at_depth(
+                far_z, image_height, image_width, cam)
+            
+            card_corners_3d = np.array([
+                [-far_width / 2, -far_height / 2, far_z],
+                [far_width / 2, -far_height / 2, far_z],
+                [near_width / 2, near_height / 2, near_z + 2 * far_z], # NO IDEA
+                [-near_width / 2, near_height / 2, near_z + 2 * far_z]
+            ], dtype=np.float32)
+            
+            print(card_corners_3d)
+        else:
+            z = self._depth_to_z(self.depth, cam)
+
+            # Calculate the 3D points of the card corners
+            card_width, card_height = self._dimension_at_depth(
+                z, image_height, image_width, cam)
+            
+            card_corners_3d = np.array([
+                [-card_width / 2, -card_height / 2, z],
+                [card_width / 2, -card_height / 2, z],
+                [card_width / 2, card_height / 2, z],
+                [-card_width / 2, card_height / 2, z]
+            ], dtype=np.float32)
 
         return card_corners_3d