diff --git a/src/robofish/io/app.py b/src/robofish/io/app.py index 3b9dcd9840555d81de4d188768bb9d42f973d6ac..a9344741f866ff3a6ad1fb4f85d4f44cef1c1161 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 0b5c4f169ff3c494370541271532b16d7ca9b3a4..cc75d522c5f364e18854ef9f5bc44ae5a7d33138 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