From 7e856352e2f6ba0d549d49b69c44111139ebdab9 Mon Sep 17 00:00:00 2001
From: Andi Gerken <andi.gerken@gmail.com>
Date: Fri, 21 Mar 2025 17:58:13 +0100
Subject: [PATCH] Added smoothing for fov orientation change to make it less
stressful to watch.
---
.../load_model_params_from_model_options.py | 44 +++++++++++++++++++
src/robofish/io/app.py | 1 +
src/robofish/io/file.py | 22 +++++++---
3 files changed, 61 insertions(+), 6 deletions(-)
create mode 100644 src/conversion_scripts/load_model_params_from_model_options.py
diff --git a/src/conversion_scripts/load_model_params_from_model_options.py b/src/conversion_scripts/load_model_params_from_model_options.py
new file mode 100644
index 0000000..fb17140
--- /dev/null
+++ b/src/conversion_scripts/load_model_params_from_model_options.py
@@ -0,0 +1,44 @@
+from pathlib import Path
+import argparse
+import robofish.io
+from tqdm import tqdm
+
+def main(file_path):
+ files = list(file_path.glob("*.hdf5")) if file_path.is_dir() else [file_path]
+
+ for f in tqdm(files):
+ with robofish.io.File(f, "r+") as iof:
+ model_options = iof.attrs["model_options"]
+ model_options = eval(model_options)["options"]
+
+ neccessary_options = ["zor", "zoo", "zoa", "fov", "additive_zone_sizes"]
+
+ # extracted_options
+ e_options = {no: model_options[no] for no in neccessary_options}
+
+ if e_options["additive_zone_sizes"]:
+ e_options["zoo"] += e_options["zor"]
+ e_options["zoa"] += e_options["zoo"]
+
+ for e in iof.entities:
+ e.attrs["category"] = "organism"
+
+ if "model" in e:
+ del e["model"]
+
+ g = e.create_group("model")
+ g.attrs["name"] = "couzin"
+
+ p = g.create_group("parameters")
+ p.attrs["zoo"] = e_options["zoo"]
+ p.attrs["zoa"] = e_options["zoa"]
+ p.attrs["zor"] = e_options["zor"]
+ p.attrs["fov"] = e_options["fov"]
+
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser(description="Convert csv file from the socoro experiments to RoboFish track format.")
+ parser.add_argument("file_path", type=str, help="Path to the csv file.")
+ args = parser.parse_args()
+ main(Path(args.file_path))
\ No newline at end of file
diff --git a/src/robofish/io/app.py b/src/robofish/io/app.py
index 9c35211..4aec636 100644
--- a/src/robofish/io/app.py
+++ b/src/robofish/io/app.py
@@ -267,6 +267,7 @@ def render(args: argparse.Namespace = None) -> None:
"render_swarm_center": False,
"highlight_switches": False,
"figsize": 10,
+ "fov_smoothing_factor": 0.8
}
for key, value in default_options.items():
diff --git a/src/robofish/io/file.py b/src/robofish/io/file.py
index cc267cb..8af3217 100644
--- a/src/robofish/io/file.py
+++ b/src/robofish/io/file.py
@@ -1026,6 +1026,7 @@ class File(h5py.File):
custom_colors: bool = None,
dpi: int = 200,
figsize: int = 10,
+ fov_smoothing_factor: float = 0.8
) -> None:
"""Render a video of the file.
@@ -1116,9 +1117,9 @@ class File(h5py.File):
entity_zones = []
for zone_size in zone_sizes.values():
if fov >= 360:
- entity_zones.append(matplotlib.patches.Circle((0, 0), zone_size, color=fish_colors[ei], alpha=0.3, fill=False))
+ entity_zones.append(matplotlib.patches.Circle((0, 0), zone_size, color=fish_colors[ei], alpha=0.2, fill=False))
else:
- entity_zones.append(matplotlib.patches.Wedge((0,0), zone_size, theta1=-fov/2, theta2=fov/2, color=fish_colors[ei], alpha=0.3, fill=False))
+ entity_zones.append(matplotlib.patches.Wedge((0,0), zone_size, theta1=-fov/2, theta2=fov/2, color=fish_colors[ei], alpha=0.2, fill=False))
zones.append(entity_zones)
zones_flat = []
@@ -1319,6 +1320,8 @@ class File(h5py.File):
pbar = tqdm(range(n_frames)) if video_path is not None else None
+ self.fov_orientations = [None] * n_entities
+
def update(frame):
output_list = []
@@ -1395,15 +1398,22 @@ class File(h5py.File):
poses_trails = entity_poses[:, max(0, file_frame - trail) : file_frame]
for i_entity in range(n_entities):
- for zone in zones[i_entity]:
-
+ new_ori = this_pose[i_entity, 2]
+ old_ori = self.fov_orientations[i_entity] if self.fov_orientations[i_entity] is not None else new_ori
+
+ # Mix in complex space to handle the wrap around
+ smoothed_ori = fov_smoothing_factor * np.exp(1j * old_ori) + (1-fov_smoothing_factor) * np.exp(1j * new_ori)
+ self.fov_orientations[i_entity] = np.angle(smoothed_ori)
+
+ ori_deg = np.rad2deg(self.fov_orientations[i_entity])
+
+ for zone in zones[i_entity]:
zone.set_center((
this_pose[i_entity, 0],
this_pose[i_entity, 1],
))
if fovs[i_entity] < 360:
- ori_deg = np.rad2deg(this_pose[i_entity, 2])
- zone.theta1 = ori_deg - fovs[i_entity] / 2
+ zone.theta1 = ori_deg - fovs[i_entity] / 2
zone.theta2 = ori_deg + fovs[i_entity] / 2
if render_swarm_center:
--
GitLab