97 lines
3.0 KiB
Python
97 lines
3.0 KiB
Python
import open3d as o3d
|
|
import numpy as np
|
|
from math import pi
|
|
class Pokemon():
|
|
def __init__(self, pokemon_type, animal):
|
|
self.type = pokemon_type
|
|
self.animal = animal
|
|
|
|
self.scalar = 8
|
|
self.x = 7
|
|
self.y = 10
|
|
self.z = 5
|
|
self.theta = 2*pi
|
|
self.accum_rotate = np.array([[1, 0, 0],[0, 1, 0],[0,0,1]])
|
|
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]
|
|
self.cloudpoints['colors'] = self.cloudpoints['colors'][:, ::-1]
|
|
|
|
# down sample
|
|
point_cloud = o3d.geometry.PointCloud()
|
|
point_cloud.points = o3d.utility.Vector3dVector(self.cloudpoints['points'])
|
|
point_cloud.colors = o3d.utility.Vector3dVector(self.cloudpoints['colors'])
|
|
|
|
downsampled_cloud = point_cloud.uniform_down_sample(every_k_points=5)
|
|
downsampled_points = np.asarray(downsampled_cloud.points)
|
|
downsampled_colors = np.asarray(downsampled_cloud.colors)
|
|
|
|
self.cloudpoints['points'] = downsampled_points
|
|
self.cloudpoints['colors'] = downsampled_colors
|
|
|
|
self.cal_position()
|
|
|
|
def walk_forward(self):
|
|
self.x += 2
|
|
self.cal_position()
|
|
|
|
def walk_backward(self):
|
|
self.x -= 2
|
|
self.cal_position()
|
|
|
|
def walk_left(self):
|
|
self.y -= 2
|
|
self.cal_position()
|
|
|
|
def walk_right(self):
|
|
self.y += 2
|
|
self.cal_position()
|
|
|
|
def rotate_left(self):
|
|
self.theta += pi/4
|
|
self.cal_position()
|
|
|
|
def rotate_right(self):
|
|
self.theta -= pi/4
|
|
self.cal_position()
|
|
|
|
|
|
|
|
def cal_position(self):
|
|
transform_matrix = np.array([
|
|
[np.cos(self.theta)*self.scalar, -np.sin(self.theta)*self.scalar,0, 0],
|
|
[0,0 , self.scalar, 0],
|
|
[np.sin(self.theta)*self.scalar, np.cos(self.theta)*self.scalar, 0, 0],
|
|
[self.x, self.y, self.z, 1]]) # rotate (x, y, z, theta)
|
|
pokemon_points = self.cloudpoints['points']
|
|
|
|
'''
|
|
transform_matrix[3, :3] = transform_matrix[3, :3] @ self.accum_rotate
|
|
tmp = transform_matrix[[0,2, 1], :3]
|
|
self.accum_rotate = self.accum_rotate @ tmp
|
|
'''
|
|
|
|
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[:, 2] = -pokemon_points[:, 2]
|
|
self.wcs_point = pokemon_points
|
|
|
|
def get_position(self):
|
|
# return self.cloudpoints['points']
|
|
return self.wcs_point
|
|
|
|
def get_colors(self):
|
|
return self.cloudpoints['colors']
|
|
|