79 lines
2.4 KiB
Python
79 lines
2.4 KiB
Python
import open3d as o3d
|
|
import numpy as np
|
|
|
|
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.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 cal_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, 1]]) # 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[:, 2] = -pokemon_points[:, 2]
|
|
self.wcs_point = pokemon_points
|
|
|
|
def get_position(self):
|
|
return self.wcs_point
|
|
|
|
def get_colors(self):
|
|
return self.cloudpoints['colors']
|