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 0000000000000000000000000000000000000000..fb17140de699ae3a8b2c0419f8ee32da59dec804 --- /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 9c352111019ec9606694fcff65299ecd86d00ba5..4aec636e080ab18bea68faf12b459a118c38315f 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 cc267cbecc21958fb51b3b8068a0b1e549be126b..8af3217ae6bf7f3f5c4230fe7319f8bf6fdb95dd 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: