From 1d8fb2df30f573e4348e8df0aab3e508945140b5 Mon Sep 17 00:00:00 2001
From: mhocke <mathis.hocke@fu-berlin.de>
Date: Mon, 14 Feb 2022 14:57:17 +0000
Subject: [PATCH] Show values of text datasets in title of render view.
---
src/robofish/io/app.py | 1 +
src/robofish/io/file.py | 19 +++++++++++++++++--
2 files changed, 18 insertions(+), 2 deletions(-)
diff --git a/src/robofish/io/app.py b/src/robofish/io/app.py
index 3b9dcd9..a934474 100644
--- a/src/robofish/io/app.py
+++ b/src/robofish/io/app.py
@@ -129,6 +129,7 @@ def render(args=None):
"view_size": 60,
"slow_view": 0.8,
"cut_frames": 0,
+ "show_text": False,
}
for key, value in default_options.items():
diff --git a/src/robofish/io/file.py b/src/robofish/io/file.py
index 0b5c4f1..cc75d52 100644
--- a/src/robofish/io/file.py
+++ b/src/robofish/io/file.py
@@ -836,6 +836,7 @@ class File(h5py.File):
"slow_view": 0.8,
"slow_zoom": 0.95,
"cut_frames": None,
+ "show_text": False,
}
options = {
@@ -851,9 +852,12 @@ class File(h5py.File):
plt.plot([], [], lw=options["linewidth"], zorder=0)[0]
for _ in range(n_entities)
]
+ categories = [entity.attrs.get("category", None) for entity in self.entities]
entity_polygons = [
- patches.Polygon(shape_vertices(options["entity_scale"]), facecolor="k")
- for _ in range(n_entities)
+ patches.Polygon(shape_vertices(options["entity_scale"]), facecolor=color)
+ for color in [
+ "gray" if category == "robot" else "k" for category in categories
+ ]
]
border_vertices = np.array(
@@ -883,6 +887,15 @@ class File(h5py.File):
min_view = np.max((np.max(start_pose, axis=0) - np.min(start_pose, axis=0))[:2])
self.view_size = np.max([options["view_size"], min_view + options["margin"]])
+ def title(file_frame: int) -> str:
+ """Search for datasets containing text for displaying it in the video"""
+ output = []
+ for e in self.entities:
+ for key, val in e.items():
+ if val.dtype == object and type(val[0]) == bytes:
+ output.append(f"{e.name}.{key}='{val[file_frame].decode()}'")
+ return ", ".join(output)
+
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])
@@ -943,6 +956,8 @@ class File(h5py.File):
self.middle_of_swarm[1] - self.view_size / 2,
self.middle_of_swarm[1] + self.view_size / 2,
)
+ if options["show_text"]:
+ ax.set_title(title(file_frame))
poses_trails = entity_poses[
:, max(0, file_frame - options["trail"]) : file_frame
--
GitLab