Skip to content
Snippets Groups Projects
Commit 719e2e4e authored by gianluv00's avatar gianluv00 :flag_eu:
Browse files

Changed return type of create_entity and create_multiple_entities to robofish.io.Entity

parent 0f7b93de
No related branches found
No related tags found
No related merge requests found
Pipeline #51361 failed
......@@ -30,7 +30,7 @@ class Entity(h5py.Group):
orientations: Iterable = None,
outlines: Iterable = None,
sampling: str = None,
):
) -> robofish.io.Entity:
poses, positions, orientations, outlines = utils.np_array(
poses, positions, orientations, outlines
)
......
......@@ -439,7 +439,7 @@ class File(h5py.File):
orientations: Iterable = None,
outlines: Iterable = None,
sampling: str = None,
) -> str:
) -> robofish.io.Entity:
"""Creates a new single entity.
Args:
......@@ -478,7 +478,7 @@ class File(h5py.File):
names: Iterable[str] = None,
outlines=None,
sampling=None,
) -> Iterable:
) -> Iterable[robofish.io.Entity]:
"""Creates multiple entities.
Args:
......@@ -496,7 +496,7 @@ class File(h5py.File):
), f"A 3 dimensional array was expected (entity, timestep, 3). There were {poses.ndim} dimensions in poses: {poses.shape}"
assert poses.shape[2] in [3, 4]
agents = poses.shape[0]
entity_names = []
entities = []
for i in range(agents):
e_name = None if names is None else names[i]
......@@ -504,7 +504,7 @@ class File(h5py.File):
outlines if outlines is None or outlines.ndim == 3 else outlines[i]
)
entity_names.append(
entities.append(
self.create_entity(
category=category,
sampling=sampling,
......@@ -513,7 +513,7 @@ class File(h5py.File):
outlines=e_outline,
)
)
return entity_names
return entities
def update_calculated_data(self, verbose=False):
changed = any([e.update_calculated_data(verbose) for e in self.entities])
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment