diff --git a/pyproject.toml b/pyproject.toml index f48d119..dd93d67 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -92,7 +92,7 @@ ignore_missing_imports = true convention = "numpy" add-select = [ "D212", # Multi-line docstring summary should start at the first line - "D402", # First line should not be the function’s “signature” + "D402", # First line should not be the functions "signature" "D417", # Missing argument descriptions in the docstring ] add-ignore = [ diff --git a/raillabel/__init__.py b/raillabel/__init__.py index 7188785..31182e0 100644 --- a/raillabel/__init__.py +++ b/raillabel/__init__.py @@ -3,7 +3,7 @@ """Devkit for working with recorded and annotated train ride data from DB.""" from importlib import metadata -from . import _util, format +from . import _util, format, stats from .exceptions import * from .filter.filter import filter from .format import Scene diff --git a/raillabel/stats/__init__.py b/raillabel/stats/__init__.py new file mode 100644 index 0000000..e1fb346 --- /dev/null +++ b/raillabel/stats/__init__.py @@ -0,0 +1,5 @@ +# Copyright DB Netz AG and contributors +# SPDX-License-Identifier: Apache-2.0 +"""Statsmodule provides statistics.""" + +from .generate_timespan import generate_timespan diff --git a/raillabel/stats/generate_timespan.py b/raillabel/stats/generate_timespan.py new file mode 100644 index 0000000..370410f --- /dev/null +++ b/raillabel/stats/generate_timespan.py @@ -0,0 +1,36 @@ +# Copyright DB Netz AG and contributors +# SPDX-License-Identifier: Apache-2.0 + +import typing as t +from decimal import Decimal + +from ..format import Scene + + +def generate_timespan(scene: Scene) -> t.Tuple[t.Optional[Decimal], t.Optional[Decimal]]: + """Return start and end timestamp of the scene. + + Parameters + ---------- + scene: raillabel.format.Scene + Scene the timespan should be based off. + + Returns + ------- + decimal.Decimal or None + Start timestamp of the scene. Is None if the scene has no frames. + decimal.Decimal or None + End timestamp of the scene. Is None if the scene has no frames. + """ + start_timestamp = None + end_timestamp = None + + for frame in scene.frames.values(): + + if start_timestamp == None or frame.timestamp < start_timestamp: + start_timestamp = frame.timestamp + + if end_timestamp == None or frame.timestamp > end_timestamp: + end_timestamp = frame.timestamp + + return (start_timestamp, end_timestamp) diff --git a/tests/test_raillabel/stats/test_generate_timespan.py b/tests/test_raillabel/stats/test_generate_timespan.py new file mode 100644 index 0000000..46b159b --- /dev/null +++ b/tests/test_raillabel/stats/test_generate_timespan.py @@ -0,0 +1,71 @@ +# Copyright DB Netz AG and contributors +# SPDX-License-Identifier: Apache-2.0 + +import os +import sys +from pathlib import Path + +import pytest + +sys.path.insert(1, str(Path(__file__).parent.parent.parent.parent)) + +import raillabel + +# == Fixtures ============================ + +@pytest.fixture +def metadata() -> raillabel.format.Metadata: + return raillabel.format.Metadata(schema_version="1.0.0") + +# == Tests ============================ + +def test_simple_timespan(metadata): + scene = raillabel.Scene( + metadata=metadata, + frames={ + 0: raillabel.format.Frame( + uid=0, + timestamp=100 + ), + 1: raillabel.format.Frame( + uid=1, + timestamp=105 + ), + 2: raillabel.format.Frame( + uid=0, + timestamp=110 + ), + } + ) + + assert raillabel.stats.generate_timespan(scene) == (100, 110) + +def test_unordered_timspan(metadata): + scene = raillabel.Scene( + metadata=metadata, + frames={ + 0: raillabel.format.Frame( + uid=0, + timestamp=110 + ), + 1: raillabel.format.Frame( + uid=1, + timestamp=100 + ), + } + ) + + assert raillabel.stats.generate_timespan(scene) == (100, 110) + +def test_empty_timespan(metadata): + scene = raillabel.Scene( + metadata=metadata, + frames={} + ) + + assert raillabel.stats.generate_timespan(scene) == (None, None) + + +if __name__ == "__main__": + os.system("clear") + pytest.main([__file__, "--disable-pytest-warnings", "--cache-clear", "-v"])