From 048f15cb08302ac081fb2d1d3ea31f9f6cfa85dd Mon Sep 17 00:00:00 2001
From: Ting-Jun Wang <levi900227@gmail.com>
Date: Sun, 17 Dec 2023 22:49:54 +0800
Subject: [PATCH] feat: render pokemon to Open3D

---
 main.py    | 20 ++++++++++++++++----
 pokemon.py | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+), 4 deletions(-)
 create mode 100644 pokemon.py

diff --git a/main.py b/main.py
index b073b8b..feb926d 100644
--- a/main.py
+++ b/main.py
@@ -1,6 +1,7 @@
 import numpy as np
 import open3d as o3d
 import cv2
+from pokemon import Pokemon
 
 
 PAPER_WCS_POINT = np.array([[0, 0, 0], [18.5, 0, 0], [0, 26, 0], [18.5, 26, 0]], dtype=np.float32)
@@ -81,7 +82,7 @@ def plot_corner_points(frame, approx):
     cv2.circle(frame, (x, y), 30, (255, 0, 0), -1)  # 在角點位置畫紅色圓圈
 
 
-def visualization(inverse_matrix_M):
+def visualization(inverse_matrix_M, pokemon):
     box_size = 0.3
     
     # initialize visualizer
@@ -112,6 +113,14 @@ def visualization(inverse_matrix_M):
     paper_triangle_set.triangles = o3d.utility.Vector3iVector(paper_triangles)
     vis.add_geometry(paper_triangle_set)
 
+
+    # Pokemon model
+    pcd = o3d.geometry.PointCloud()
+    pcd.points = o3d.utility.Vector3dVector(pokemon.get_position())
+    pcd.colors = o3d.utility.Vector3dVector(pokemon.get_colors())
+    vis.add_geometry(pcd)
+
+
     
     # 下面是畫相機
     # 找 金字塔端點
@@ -182,10 +191,12 @@ def visualization(inverse_matrix_M):
 if __name__ == '__main__':
 
     # 設定連接到 Android 手機的相機
-    cap = cv2.VideoCapture(0)  # 0 表示第一個相機（通常是後置相機），若是前置相機，可以使用 1
-    # cap = cv2.VideoCapture('demo1.mp4')
+    # cap = cv2.VideoCapture(0)  # 0 表示第一個相機（通常是後置相機），若是前置相機，可以使用 1
+    cap = cv2.VideoCapture('demo1.mp4')
 
     previous_approx = []
+
+    pokemon = Pokemon('water', 'cat')
     
     rotation_vectors = []
     translation_vectors = []
@@ -274,7 +285,8 @@ if __name__ == '__main__':
         M = np.concatenate((rt, [[0, 0, 0, 1]]), axis=0)
         M_inv.append(np.linalg.inv(M))
 
-    visualization(M_inv)
+    # load Pokemon points
+    visualization(M_inv, pokemon)
 
 
     
diff --git a/pokemon.py b/pokemon.py
new file mode 100644
index 0000000..6d9a1a9
--- /dev/null
+++ b/pokemon.py
@@ -0,0 +1,46 @@
+
+import numpy as np
+
+class Pokemon():
+    def __init__(self, pokemon_type, animal):
+        self.type = pokemon_type
+        self.animal = animal
+
+        self.scalar = 10
+        
+        self.x = 8
+        self.y = 13
+        self.z = 8
+        
+        self.cloudpoints = {
+            'points': None,
+            'colors': None,
+        }
+
+        self.load_model()
+
+    def load_model(self):
+        self.cloudpoints['points'] = np.load('models/{}_{}/0/points.npy'.format(self.type, self.animal))
+        self.cloudpoints['colors'] = np.load('models/{}_{}/0/colors.npy'.format(self.type, self.animal))
+
+        self.cloudpoints['colors'] = self.cloudpoints['colors'] / 255
+        self.cloudpoints['colors'] = self.cloudpoints['colors'][:, :3]
+
+    def get_position(self):
+        transform_matrix = np.array([
+            [self.scalar, 0, 0, 0],
+            [0, 0, self.scalar, 0],
+            [0, self.scalar, 0, 0],
+            [self.x, self.y, self.z, 0]])             # shift (前後, 左右, 高度)
+
+        pokemon_points = self.cloudpoints['points']
+
+        ones = np.ones((pokemon_points.shape[0], 1))
+        pokemon_points = np.concatenate([pokemon_points, ones], axis=1)
+        pokemon_points = pokemon_points @ transform_matrix
+        pokemon_points = pokemon_points[:, :3]
+
+        return pokemon_points
+    
+    def get_colors(self):
+        return self.cloudpoints['colors']