Skip to content

Commit

Permalink
feat: implement ExcludeAnnotationTypeFilter
Browse files Browse the repository at this point in the history
  • Loading branch information
unexcellent committed Nov 16, 2024
1 parent f41ce29 commit fd3cc1f
Show file tree
Hide file tree
Showing 3 changed files with 56 additions and 0 deletions.
2 changes: 2 additions & 0 deletions raillabel/filter/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@

from .end_time_filter import EndTimeFilter
from .exclude_annotation_id_filter import ExcludeAnnotationIdFilter
from .exclude_annotation_type_filter import ExcludeAnnotationTypeFilter
from .exclude_frame_id_filter import ExcludeFrameIdFilter
from .filter import filter_
from .include_annotation_id_filter import IncludeAnnotationIdFilter
Expand All @@ -20,4 +21,5 @@
"IncludeAnnotationIdFilter",
"ExcludeAnnotationIdFilter",
"IncludeAnnotationTypeFilter",
"ExcludeAnnotationTypeFilter",
]
46 changes: 46 additions & 0 deletions raillabel/filter/exclude_annotation_type_filter.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
# Copyright DB InfraGO AG and contributors
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

from dataclasses import dataclass
from typing import Literal
from uuid import UUID

from raillabel.format import Bbox, Cuboid, Poly2d, Poly3d, Seg3d

from ._filter_abc import _AnnotationLevelFilter


@dataclass
class ExcludeAnnotationTypeFilter(_AnnotationLevelFilter):
"""Filter out all annotations in the scene, that do have the type (like bbox or cuboid)."""

annotation_types: (
set[Literal["bbox", "cuboid", "poly2d", "poly3d", "seg3d"]]
| list[Literal["bbox", "cuboid", "poly2d", "poly3d", "seg3d"]]
)

def passes_filter(self, _: UUID, annotation: Bbox | Cuboid | Poly2d | Poly3d | Seg3d) -> bool:
"""Assess if an annotation passes this filter."""
annotation_type_str = None

if isinstance(annotation, Bbox):
annotation_type_str = "bbox"

elif isinstance(annotation, Cuboid):
annotation_type_str = "cuboid"

elif isinstance(annotation, Poly2d):
annotation_type_str = "poly2d"

elif isinstance(annotation, Poly3d):
annotation_type_str = "poly3d"

elif isinstance(annotation, Seg3d):
annotation_type_str = "seg3d"

else:
raise TypeError

return annotation_type_str not in self.annotation_types
8 changes: 8 additions & 0 deletions tests/filter/test_filter.py
Original file line number Diff line number Diff line change
Expand Up @@ -100,5 +100,13 @@ def test_include_annotation_type():
assert actual == SceneBuilder.empty().add_bbox().result


def test_exclude_annotation_type():
scene = SceneBuilder.empty().add_bbox().add_cuboid().result
filters = [raillabel.filter.ExcludeAnnotationTypeFilter(["cuboid"])]

actual = raillabel.filter.filter_(scene, filters)
assert actual == SceneBuilder.empty().add_bbox().result


if __name__ == "__main__":
pytest.main([__file__, "-vv"])

0 comments on commit fd3cc1f

Please sign in to comment.