Skip to content
Snippets Groups Projects
Commit e9175f62 authored by Andi Gerken's avatar Andi Gerken
Browse files
parents f4d15b8e ee5990df
Branches
Tags 0.2.12
No related merge requests found
Pipeline #48751 passed
......@@ -137,6 +137,8 @@ def render(args=None):
"slow_view": 0.8,
"cut_frames": 0,
"show_text": False,
"render_goals": False,
"render_targets": False,
}
for key, value in default_options.items():
......
......@@ -22,7 +22,7 @@ import h5py
import numpy as np
import logging
from typing import Iterable, Union, Tuple, List
from typing import Iterable, Union, Tuple, List, Optional
from pathlib import Path
import shutil
import datetime
......@@ -844,6 +844,8 @@ class File(h5py.File):
"slow_zoom": 0.95,
"cut_frames": None,
"show_text": False,
"render_goals": False,
"render_targets": False,
}
options = {
......@@ -859,6 +861,10 @@ class File(h5py.File):
plt.plot([], [], lw=options["linewidth"], zorder=0)[0]
for _ in range(n_entities)
]
points = [
plt.scatter([], [], marker="x", color="k"),
plt.plot([], [], linestyle="dotted", alpha=0.5, color="k", zorder=0)[0],
]
categories = [entity.attrs.get("category", None) for entity in self.entities]
entity_polygons = [
patches.Polygon(shape_vertices(options["entity_scale"]), facecolor=color)
......@@ -903,6 +909,28 @@ class File(h5py.File):
output.append(f"{e.name}.{key}='{val[file_frame].decode()}'")
return ", ".join(output)
def get_goal(file_frame: int) -> Optional[np.ndarray]:
"""Return current goal of robot, if robot exists and has a goal."""
goal = None
if "robot" in categories:
robot = self.entities[categories.index("robot")]
try:
goal = robot["goals"][file_frame]
except KeyError:
pass
if goal is not None and np.isnan(goal).any():
goal = None
return goal
def get_target(file_frame: int) -> Tuple[List, List]:
"""Return line points from robot to target"""
if "robot" in categories:
robot = self.entities[categories.index("robot")]
rpos = robot["positions"][file_frame]
target = robot["targets"][file_frame]
return [rpos[0], target[0]], [rpos[1], target[1]]
return [], []
def init():
ax.set_xlim(-0.5 * self.world_size[0], 0.5 * self.world_size[0])
ax.set_ylim(-0.5 * self.world_size[1], 0.5 * self.world_size[1])
......@@ -914,7 +942,7 @@ class File(h5py.File):
for e_poly in entity_polygons:
ax.add_patch(e_poly)
ax.add_patch(border)
return lines + entity_polygons + [border]
return lines + entity_polygons + [border] + points
n_frames = self.entity_poses.shape[1]
if options["cut_frames"] > 0:
......@@ -970,6 +998,14 @@ class File(h5py.File):
if options["show_text"]:
ax.set_title(title(file_frame))
if options["render_goals"]:
goal = get_goal(file_frame)
if goal is not None:
points[0].set_offsets(goal)
if options["render_targets"]:
points[1].set_data(get_target(file_frame))
poses_trails = entity_poses[
:, max(0, file_frame - options["trail"]) : file_frame
]
......@@ -989,7 +1025,7 @@ class File(h5py.File):
raise Exception(
f"Frame is bigger than n_frames {file_frame} of {n_frames}"
)
return lines + entity_polygons + [border]
return lines + entity_polygons + [border] + points
print(f"Preparing to render n_frames: {n_frames}")
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment