PokeFusion/pokemon.py

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]
# 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.cloudpoints['points']
return self.wcs_point
def get_colors(self):
return self.cloudpoints['colors']