def main():
for num_points in [100]:
models_path = "./data/1kabc/simple/val"
optimal_numbers, losses = [], []
for model_name in tqdm(sorted(os.listdir(models_path))):
model_path = os.path.join(models_path, model_name)
# model_path = "./data/1kabc/simple/train/00070090_73b2f35a88394199b6fd1ab8_003.obj"
model = Model(model_path)
model.generate_view_points(num_points)
optimal, loss = find_greedy_optimal(model, do_rec=True)
optimal_numbers.append(len(optimal))
losses.append(loss)
print("Model: ", model_path, "Optimal number: ",
np.mean(optimal_numbers), "Loss: ", np.mean(losses))
def reset(self):
"""
Reset the environment for new episode.
Randomly (or not) generate CAD model for this episode.
"""
if self.models_path is not None:
self.model_path = os.path.join(self.models_path,
random.sample(os.listdir(self.models_path), 1)[0])
self.model = Model(self.model_path, resolution_image=self.image_size)
self.model.generate_view_points(self.number_of_view_points)
if self.illustrate:
self.model.illustrate().display()
self.plot = k3d.plot()
self.plot.display()
init_action = self.action_space.sample()
observation = self.model.get_observation(init_action)
return observation, init_action
def create_axis(dimensions, size, color, **kw):
return Model(
dimensions,
[
Primitive(
Line(
Vector.get_zero_vector(dimensions),
Vector(*(size if i == d else 0 for i in range(dimensions)))
),
color,
arrow=LAST)
for d in range(dimensions)
],
**kw)
def create_cube(size, dimensions, color="black", **kw):
lines = []
for d in range(dimensions**2):
point = Vector(*(size * (d & 2**i == 2**i) for i in range(dimensions)))
for n in range(dimensions):
if point.coordinates[n] == 0:
lines.append(
Line(
point,
Vector(*(point.coordinates[i] if i != n else size
for i in range(dimensions)))))
return Model(dimensions, [Primitive(l, color) for l in lines], **kw)
class Environment(gym.Env):
def __init__(self,
models_path=None,
model_path=None,
number_of_view_points=100,
similarity_threshold=0.95,
image_size=512,
illustrate=False):
super().__init__()
self.model_path = model_path
self.models_path = models_path
self.number_of_view_points = number_of_view_points
self.image_size = image_size
self.illustrate = illustrate
self.action_space = spaces.Discrete(number_of_view_points)
self.observation_space = spaces.Dict({
'depth_map':
spaces.Box(-np.inf,
np.inf, (image_size, image_size),
dtype=np.float32),
})
self._similarity_threshold = similarity_threshold
self._reconstruction_depth = 10
self.model = None
self.plot = None
def reset(self):
"""
Reset the environment for new episode.
Randomly (or not) generate CAD model for this episode.
"""
if self.models_path is not None:
self.model_path = os.path.join(
self.models_path,
random.sample(os.listdir(self.models_path), 1)[0])
self.model = Model(self.model_path, resolution_image=self.image_size)
self.model.generate_view_points(self.number_of_view_points)
if self.illustrate:
self.model.illustrate().display()
self.plot = k3d.plot()
self.plot.display()
init_action = self.action_space.sample()
observation = self.model.get_observation(init_action)
return observation, init_action
def step(self, action):
"""
Get new observation from current position (action), count step reward, decide whether to stop.
Args:
action: int
return:
next_state: List[List[List[int, int, int]]]
reward: float
done: bool
info: Tuple
"""
assert self.action_space.contains(action)
observation = self.model.get_observation(action)
reward = self.step_reward(observation)
done = reward >= self._similarity_threshold
return observation, reward, done, {}
def render(self, action, observation, plot=None):
if plot is None:
plot = self.plot
plot = illustrate_points([self.model.get_point(action)],
size=0.5,
plot=plot)
plot = observation.illustrate(plot, size=0.03)
return plot
def step_reward(self, observation):
# THINK ABOUT yet another reward
return self.model.observation_similarity(observation)
if self.illustrate:
illustrate_mesh(vertices, faces).display()
return reward
def _get_mesh(self, observation):
faces, vertices = poisson_reconstruction(
observation.points,
observation.normals,
depth=self._reconstruction_depth)
return vertices, faces
def final_reward(self, observation):
vertices, faces = self._get_mesh(observation)
reward = self.model.surface_similarity(vertices, faces)
return reward