Repository: una-auxme/paf
Issue category: Feature Request

Issue:
[Feature]: Intermediate layer: integrate shapely into Shape2D
### Feature Description

Shapely description: Manipulation and analysis of geometric objects in the Cartesian plane

Might be a basic building block to "easily" develop algorithms that extract useful data from the intermediate layer

### Definition of Done

- [x] Create shapely shape from Shape2D
- [x] Use the shapely.STRtree for accelerating the intersection checks

Work related to the following functions can be found in #599 , #601 :
- Functions for calculating intersection, union, difference between shapes
- Calculate IoU between shapes

### Testability

- [x] Units tests that roughly evaluate the shape creation functions

### Dependencies

- [x] Intermediate layer


Hints:

Patch:
diff --git a/code/mapping/ext_modules/mapping_common/entity.py b/code/mapping/ext_modules/mapping_common/entity.py
index 5c84f42f..19a59dd6 100644
--- a/code/mapping/ext_modules/mapping_common/entity.py
+++ b/code/mapping/ext_modules/mapping_common/entity.py
@@ -2,6 +2,8 @@
 from enum import Enum
 from dataclasses import dataclass, field
 
+import shapely
+
 from uuid import UUID, uuid4
 from genpy.rostime import Time, Duration
 from std_msgs.msg import Header
@@ -369,6 +371,9 @@ def to_marker(self) -> Marker:
 
         return m
 
+    def to_shapely(self) -> "ShapelyEntity":
+        return ShapelyEntity(self, self.shape.to_shapely(self.transform))
+
 
 @dataclass(init=False)
 class Car(Entity):
@@ -477,3 +482,16 @@ def __init__(self, **kwargs):
 for t in _entity_supported_classes:
     t_name = t.__name__.lower()
     _entity_supported_classes_dict[t_name] = t
+
+
+@dataclass
+class ShapelyEntity:
+    """A Container containing both an entity and a shapely.Polygon
+    based on the entity's shape and transform
+
+    **IMPORTANT** If the entity is modified, the Polygon will
+    not automatically update itself and will contain outdated information.
+    """
+
+    entity: Entity
+    poly: shapely.Polygon
diff --git a/code/mapping/ext_modules/mapping_common/map.py b/code/mapping/ext_modules/mapping_common/map.py
index a281c7e5..2d6dc6f5 100644
--- a/code/mapping/ext_modules/mapping_common/map.py
+++ b/code/mapping/ext_modules/mapping_common/map.py
@@ -1,10 +1,15 @@
 from dataclasses import dataclass, field
-from typing import List, Optional
+from typing import List, Optional, Callable, Literal, Tuple
+
+import shapely
+from shapely import STRtree
+import numpy as np
+import numpy.typing as npt
 
 from genpy.rostime import Time
 from std_msgs.msg import Header
 
-from mapping_common.entity import Entity
+from mapping_common.entity import Entity, FlagFilter, ShapelyEntity
 
 from mapping import msg
 
@@ -67,6 +72,49 @@ def entities_without_hero(self) -> List[Entity]:
             return self.entities[1:]
         return self.entities
 
+    def build_tree(
+        self,
+        f: Optional[FlagFilter] = None,
+        filter_fn: Optional[Callable[[Entity], bool]] = None,
+    ) -> "MapTree":
+        """Creates a filtered MapTree
+
+        **IMPORTANT**: The map
+        MUST NOT BE MODIFIED WHILE USING THE TREE,
+        otherwise results will be invalid or crash
+
+         Useful for for quickly calculating which entities of a map are
+        the nearest or (intersect, touch, etc.) with a given geometry
+
+        Args:
+            f (Optional[FlagFilter], optional): Filtering with FlagFilter.
+                Defaults to None.
+            filter_fn (Optional[Callable[[Entity], bool]], optional):
+                Filtering with function/lambda. Defaults to None.
+
+        Returns:
+            MapTree
+        """
+        return MapTree(self, f, filter_fn)
+
+    def filtered(
+        self,
+        f: Optional[FlagFilter] = None,
+        filter_fn: Optional[Callable[[Entity], bool]] = None,
+    ) -> List[Entity]:
+        """Filters self.entities
+
+        Args:
+            f (Optional[FlagFilter], optional): Filtering with FlagFilter.
+                Defaults to None.
+            filter_fn (Optional[Callable[[Entity], bool]], optional):
+                Filtering with function/lambda. Defaults to None.
+
+        Returns:
+            List[Entity]: List of entities both filters were matching for
+        """
+        return [e for e in self.entities if _entity_matches_filter(e, f, filter_fn)]
+
     @staticmethod
     def from_ros_msg(m: msg.Map) -> "Map":
         entities = list(map(lambda e: Entity.from_ros_msg(e), m.entities))
@@ -76,3 +124,167 @@ def to_ros_msg(self) -> msg.Map:
         entities = list(map(lambda e: e.to_ros_msg(), self.entities))
         header = Header(stamp=self.timestamp)
         return msg.Map(header=header, entities=entities)
+
+
+@dataclass(init=False)
+class MapTree:
+    """An acceleration structure around the shapely.STRtree
+
+    **IMPORTANT**: The map this tree was created with
+    MUST NOT BE MODIFIED WHILE USING THE TREE,
+    otherwise results will be invalid or crash
+
+    Useful for for quickly calculating which entities of a map are
+    the nearest or (intersect, touch, etc.) with a given geometry
+
+    The map's entities can optionally be filtered upon tree creation
+    """
+
+    _str_tree: STRtree
+    filtered_entities: List[ShapelyEntity]
+    """Only the entities of this tree that weren't filtered out from the map.
+
+    Also includes their shapely.Polygon
+    """
+    map: Map
+    """The map this tree was created with
+    """
+
+    def __init__(
+        self,
+        map: Map,
+        f: Optional[FlagFilter] = None,
+        filter_fn: Optional[Callable[[Entity], bool]] = None,
+    ):
+        """Creates a shapely.STRtree based on the given map and filtering
+
+        Both filtering methods can be combined.
+        Both filters need to match for the entity to match.
+
+        Args:
+            map (Map)
+            f (Optional[FlagFilter], optional): Filtering with FlagFilter.
+                Defaults to None.
+            filter_fn (Optional[Callable[[Entity], bool]], optional):
+                Filtering with function/lambda. Defaults to None.
+        """
+        self.map = map
+        self.filtered_entities = [e.to_shapely() for e in map.filtered(f, filter_fn)]
+        self._str_tree = STRtree(geoms=[e.poly for e in self.filtered_entities])
+
+    def _idxs_to_entity(self, idxs: npt.NDArray) -> List[ShapelyEntity]:
+        return [self.filtered_entities[i] for i in idxs]
+
+    def nearest(self, geo: shapely.Geometry) -> Optional[ShapelyEntity]:
+        """Returns the nearest Entity inside the tree based on geo
+
+        More information here:
+        https://shapely.readthedocs.io/en/stable/strtree.html#shapely.STRtree.nearest
+
+        Args:
+            geo (shapely.Geometry): Geometry to calculate the nearest entity to
+
+        Returns:
+            Optional[ShapelyEntity]: If the tree is empty, will return None.
+                Otherwise will return the nearest Entity
+        """
+        idx = self._str_tree.nearest(geo)
+        if idx is None:
+            return None
+        return self.filtered_entities[idx]
+
+    def query(
+        self,
+        geo: shapely.Geometry,
+        predicate: Optional[
+            Literal[
+                "intersects",
+                "within",
+                "contains",
+                "overlaps",
+                "crosses",
+                "touches",
+                "covers",
+                "covered_by",
+                "contains_properly",
+                "dwithin",
+            ]
+        ] = None,
+        distance: Optional[float] = None,
+    ) -> List[ShapelyEntity]:
+        """Calculates which entities interact with *geo*
+
+        More information here:
+        https://shapely.readthedocs.io/en/stable/strtree.html#shapely.STRtree.query
+
+        Args:
+            geo (shapely.Geometry): The geometry to query with
+            predicate (Optional[ Literal[ &quot;intersects&quot;, &quot;within&quot;,
+            &quot;contains&quot;, &quot;overlaps&quot;, &quot;crosses&quot;,
+            &quot;touches&quot;, &quot;covers&quot;, &quot;covered_by&quot;,
+            &quot;contains_properly&quot;, &quot;dwithin&quot;, ] ], optional):
+            Which interaction to filter for. Defaults to None.
+            distance (Optional[float], optional):
+            Must only be set for the &quot;dwithin&quot; predicate
+            and controls its distance. Defaults to None.
+
+        Returns:
+            List[ShapelyEntity]: The List of queried entities in the tree
+        """
+        idxs: npt.NDArray[np.int64] = self._str_tree.query(
+            geo, predicate=predicate, distance=distance
+        )
+        return self._idxs_to_entity(idxs)
+
+    def query_nearest(
+        self,
+        geo: shapely.Geometry,
+        max_distance: Optional[float] = None,
+        exclusive: bool = False,
+        all_matches: bool = True,
+    ) -> List[Tuple[ShapelyEntity, float]]:
+        """Queries the distance from *geo* to its nearest entities in the tree
+
+        More information here:
+        https://shapely.readthedocs.io/en/stable/strtree.html#shapely.STRtree.query_nearest
+
+        Args:
+            geo (shapely.Geometry): The geometry to query with
+            max_distance (Optional[float], optional): Maximum distance for the query.
+                Defaults to None.
+            exclusive (bool, optional): If True, ignores entities
+            with a shape equal to geo. Defaults to False.
+            all_matches (bool, optional): If True, all equidistant and intersected
+            geometries will be returned. If False only the nearest. Defaults to True.
+
+        Returns:
+            List[Tuple[ShapelyEntity, float]]: A List of Tuples.
+            Each contains a queried entity and its distance to geo
+        """
+        query: Tuple[npt.NDArray[np.int64], npt.NDArray[np.float64]] = (
+            self._str_tree.query_nearest(
+                geo,
+                max_distance=max_distance,
+                return_distance=True,
+                exclusive=exclusive,
+                all_matches=all_matches,
+            )
+        )
+        result: List[Tuple[ShapelyEntity, float]] = []
+        for idx, distance in zip(query[0], query[1]):
+            result.append((self.filtered_entities[idx], distance))
+        return result
+
+
+def _entity_matches_filter(
+    e: Entity,
+    f: Optional[FlagFilter] = None,
+    filter_fn: Optional[Callable[[Entity], bool]] = None,
+) -> bool:
+    if f is not None:
+        if not e.matches_filter(f):
+            return False
+    if filter_fn is not None:
+        if not filter_fn(e):
+            return False
+    return True
diff --git a/code/mapping/ext_modules/mapping_common/shape.pxd b/code/mapping/ext_modules/mapping_common/shape.pxd
index e69de29b..d22ddda9 100644
--- a/code/mapping/ext_modules/mapping_common/shape.pxd
+++ b/code/mapping/ext_modules/mapping_common/shape.pxd
@@ -0,0 +1,1 @@
+from .transform cimport Transform2D, Point2, Vector2
diff --git a/code/mapping/ext_modules/mapping_common/shape.py b/code/mapping/ext_modules/mapping_common/shape.py
index 18bf323e..0a367910 100644
--- a/code/mapping/ext_modules/mapping_common/shape.py
+++ b/code/mapping/ext_modules/mapping_common/shape.py
@@ -1,12 +1,19 @@
 from dataclasses import dataclass
-import rospy
 from typing import List, Optional
-from .transform import Transform2D, Point2
+
+import shapely
+import math
+
+import rospy
 from mapping import msg
 from tf.transformations import quaternion_from_euler
 from visualization_msgs.msg import Marker
 from geometry_msgs.msg import Point
 
+from .transform import Transform2D, Point2, Vector2
+
+CIRCLE_APPROXIMATION_LENGTH = 0.5
+
 
 @dataclass
 class Shape2D:
@@ -74,6 +81,17 @@ def to_marker(self, transform: Transform2D) -> Marker:
         m.scale.z = 1.0
         return m
 
+    def to_shapely(self, transform: Transform2D) -> shapely.Polygon:
+        """Creates a shapely.Polygon based on this shape
+
+        Args:
+            transform (Transform2D): Transforms the resulting Polygon
+
+        Returns:
+            Polygon
+        """
+        raise NotImplementedError
+
 
 @dataclass(init=False)
 class Rectangle(Shape2D):
@@ -118,6 +136,20 @@ def to_marker(self, transform: Transform2D) -> Marker:
 
         return m
 
+    def to_shapely(self, transform: Transform2D) -> shapely.Polygon:
+        shape_transform: Transform2D = transform * self.offset
+        half_length = self.length / 2.0
+        half_width = self.width / 2.0
+        # LeftBack, RightBack, RightFront, LeftFront
+        corners = [
+            Point2.new(-half_length, half_width),
+            Point2.new(-half_length, -half_width),
+            Point2.new(half_length, -half_width),
+            Point2.new(half_length, half_width),
+        ]
+        corners = [(shape_transform * p).to_shapely() for p in corners]
+        return shapely.Polygon(corners)
+
 
 @dataclass(init=False)
 class Circle(Shape2D):
@@ -152,12 +184,23 @@ def to_marker(self, transform: Transform2D) -> Marker:
 
         return m
 
+    def to_shapely(self, transform: Transform2D) -> shapely.Polygon:
+        shape_transform: Transform2D = transform * self.offset
+        p: Point2 = Point2.from_vector(shape_transform.translation())
+
+        outline_length = 2 * self.radius * math.pi
+        segments = outline_length / CIRCLE_APPROXIMATION_LENGTH
+        quad_segs: int = max(2, int(segments / 4))
+        return shapely.buffer(
+            p.to_shapely(), self.radius, quad_segs=quad_segs, cap_style="round"
+        )
+
 
 @dataclass(init=False)
 class Polygon(Shape2D):
     """Polygon defined by a list of Point2 objects."""
 
-    # The points attribute does not need a redundant point for start and end
+    # The points attribute does not have a redundant point for start and end
     points: List[Point2]
 
     def __init__(self, points: List[Point2], offset: Optional[Transform2D] = None):
@@ -211,6 +254,39 @@ def to_marker(self, transform: Transform2D) -> Marker:
 
         return m
 
+    def to_shapely(self, transform: Transform2D) -> shapely.Polygon:
+        shape_transform: Transform2D = transform * self.offset
+        poly_points = [(shape_transform * p).to_shapely() for p in self.points]
+        return shapely.Polygon(poly_points)
+
+    @staticmethod
+    def from_shapely(poly: shapely.Polygon, make_centered: bool = False) -> "Polygon":
+        """Creates a Polygon from a shapely.Polygon
+
+        If make_centered is True, the zero-point of the resulting polygon points
+        will be the centroid of poly.
+        And the offset will be the translation of the centroid.
+
+        If make_centered is False, the points will match the points of poly
+        and no offset will be applied.
+
+        Returns:
+            Polygon
+        """
+        coords = poly.exterior.coords
+        assert len(coords) >= 3, "Polygon requires at least 3 points."
+        coords = coords[:-1]
+        if make_centered:
+            center = poly.centroid
+            c_vec = Vector2.new(center.x, center.y)
+            transform = Transform2D.new_translation(c_vec)
+        else:
+            c_vec = Vector2.zero()
+            transform = Transform2D.identity()
+
+        points = [Point2.new(x, y) - c_vec for (x, y) in coords]
+        return Polygon(points, offset=transform)
+
 
 _shape_supported_classes = [Rectangle, Circle, Polygon]
 _shape_supported_classes_dict = {}
diff --git a/code/mapping/ext_modules/mapping_common/transform.py b/code/mapping/ext_modules/mapping_common/transform.py
index f1093689..5a29d6af 100644
--- a/code/mapping/ext_modules/mapping_common/transform.py
+++ b/code/mapping/ext_modules/mapping_common/transform.py
@@ -3,6 +3,7 @@
 import numpy as np
 import numpy.typing as npt
 import math
+import shapely
 
 from geometry_msgs import msg as geometry_msgs
 from mapping import msg
@@ -66,6 +67,9 @@ def from_ros_msg(m: geometry_msgs.Point) -> "Point2":
     def to_ros_msg(self) -> geometry_msgs.Point:
         return geometry_msgs.Point(x=self.x(), y=self.y(), z=0.0)
 
+    def to_shapely(self) -> shapely.Point:
+        return shapely.Point(self._matrix[:2])
+
     def __add__(self, other):
         if isinstance(other, Vector2):
             matrix = self._matrix + other._matrix
@@ -75,6 +79,15 @@ def __add__(self, other):
             f"Unsupported operand types for +: '{type(self)}' and '{type(other)}'"
         )
 
+    def __sub__(self, other):
+        if isinstance(other, Vector2):
+            matrix = self._matrix - other._matrix
+            matrix[2] = 1.0
+            return Point2(matrix)
+        raise TypeError(
+            f"Unsupported operand types for -: '{type(self)}' and '{type(other)}'"
+        )
+
 
 @dataclass(init=False, eq=False)
 class Vector2(_Coord2):
@@ -189,10 +202,28 @@ def __add__(self, other):
             matrix = self._matrix + other._matrix
             matrix[2] = 1.0
             return Vector2(matrix)
+        if isinstance(other, Point2):
+            matrix = self._matrix + other._matrix
+            matrix[2] = 1.0
+            return Point2(matrix)
         raise TypeError(
             f"Unsupported operand types for +: '{type(self)}' and '{type(other)}'"
         )
 
+    def __sub__(self, other):
+        if isinstance(other, Vector2):
+            matrix = self._matrix - other._matrix
+            matrix[2] = 1.0
+            return Vector2(matrix)
+        raise TypeError(
+            f"Unsupported operand types for -: '{type(self)}' and '{type(other)}'"
+        )
+
+    def __neg__(self) -> "Vector2":
+        matrix = -self._matrix
+        matrix[2] = 1.0
+        return Vector2(matrix)
+
 
 @dataclass(init=False, eq=False)
 class Transform2D:
diff --git a/code/requirements.txt b/code/requirements.txt
index 83bcf9bb..07cb97a0 100644
--- a/code/requirements.txt
+++ b/code/requirements.txt
@@ -19,3 +19,4 @@ pandas==2.0.3
 debugpy==1.8.7
 pyopenssl==24.3.0
 cython==3.0.11
+shapely==2.0.6
\ No newline at end of file


Test patch:
diff --git a/code/mapping/tests/mapping_common/test_entity.py b/code/mapping/tests/mapping_common/test_entity.py
index 7c1f1aa7..c308c911 100644
--- a/code/mapping/tests/mapping_common/test_entity.py
+++ b/code/mapping/tests/mapping_common/test_entity.py
@@ -10,6 +10,7 @@ def get_transform():
 
 
 def get_car():
+    flags = entity.Flags(is_collider=True)
     return entity.Car(
         indicator=entity.Car.IndicatorState.LEFT,
         confidence=1.0,
@@ -17,6 +18,7 @@ def get_car():
         shape=get_shape(),
         transform=get_transform(),
         motion=get_motion(),
+        flags=flags,
     )
 
 
diff --git a/code/mapping/tests/mapping_common/test_shape.py b/code/mapping/tests/mapping_common/test_shape.py
index b53ab889..feafcc31 100644
--- a/code/mapping/tests/mapping_common/test_shape.py
+++ b/code/mapping/tests/mapping_common/test_shape.py
@@ -17,13 +17,17 @@ def test_circle_conversion():
     assert s == s_conv
 
 
-def test_polygon_conversion():
+def get_polygon():
     data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
     p_list = []
     for i in range(0, 6, 2):
         p_list.append(Point2.new(data[i], data[i + 1]))
 
-    polygon = shape.Polygon(p_list)
+    return shape.Polygon(p_list)
+
+
+def test_polygon_conversion():
+    polygon = get_polygon()
     msg = polygon.to_ros_msg()
     polygon_conv = shape.Shape2D.from_ros_msg(msg)
     assert polygon == polygon_conv
diff --git a/code/mapping/tests/mapping_common/test_shapely.py b/code/mapping/tests/mapping_common/test_shapely.py
new file mode 100644
index 00000000..ec51a6c1
--- /dev/null
+++ b/code/mapping/tests/mapping_common/test_shapely.py
@@ -0,0 +1,121 @@
+from typing import List
+
+from mapping_common import shape, entity
+from mapping_common.map import Map
+from mapping_common.transform import Transform2D, Vector2
+
+import test_entity
+import test_shape
+
+
+def test_circle_shapely():
+    offset_transl = Vector2.new(1.0, 0.0)
+    offset = Transform2D.new_translation(offset_transl)
+    s = shape.Circle(3.0, offset)
+    poly = s.to_shapely(Transform2D.identity())
+    bounds = poly.bounds
+    assert bounds == (-2.0, -3.0, 4.0, 3.0)
+
+
+def test_rectangle_shapely():
+    s = shape.Rectangle(4.0, 1.5)
+    poly = s.to_shapely(Transform2D.identity())
+    bounds = poly.bounds
+    assert bounds == (-2.0, -0.75, 2.0, 0.75)
+
+
+def test_polygon_to_shapely():
+    s = test_shape.get_polygon()
+    poly = s.to_shapely(Transform2D.identity())
+    bounds = poly.bounds
+    assert bounds == (1.0, 2.0, 5.0, 6.0)
+
+
+def test_polygon_from_shapely():
+    s = test_shape.get_polygon()
+    poly = s.to_shapely(Transform2D.identity())
+    s_conv = shape.Polygon.from_shapely(poly)
+
+    assert s == s_conv
+
+
+def test_polygon_to_from_shapely_centered():
+    s = test_shape.get_polygon()
+    poly = s.to_shapely(Transform2D.identity())
+    s_conv = shape.Polygon.from_shapely(poly, make_centered=True)
+    poly_conv = s_conv.to_shapely(Transform2D.identity())
+
+    assert poly.equals_exact(poly_conv, tolerance=0.0)
+    assert s_conv.offset.translation() == Vector2.new(3.0, 4.0)
+
+
+def test_car_shapely():
+    car = test_entity.get_car()
+    transl = Vector2.new(2.0, -3.0)
+    transform = Transform2D.new_translation(transl)
+    car.transform = transform
+    s = car.to_shapely()
+    bounds = s.poly.bounds
+    assert bounds == (1.5, -4.0, 2.5, -2.0)
+
+
+def get_test_entities() -> List[entity.Entity]:
+    translations = [
+        Vector2.new(2.0, -3.0),
+        Vector2.new(1.0, 5.0),
+        Vector2.new(2.0, 1.0),
+    ]
+    entities = []
+    for t in translations:
+        car = test_entity.get_car()
+        transform = Transform2D.new_translation(t)
+        car.transform = transform
+        entities.append(car)
+
+    lane = entity.Lanemarking(
+        style=entity.Lanemarking.Style.SOLID,
+        confidence=1.0,
+        priority=1.0,
+        shape=shape.Rectangle(4.0, 0.2),
+        transform=Transform2D.identity(),
+        flags=entity.Flags(is_lanemark=True),
+    )
+    entities.append(lane)
+    return entities
+
+
+def test_map_tree_nearest():
+    entities = get_test_entities()
+
+    map = Map(entities=entities)
+    tree = map.build_tree(f=entity.FlagFilter(is_collider=True))
+    test_shape = shape.Circle(0.5).to_shapely(Transform2D.identity())
+    nearest = tree.nearest(test_shape)
+
+    assert nearest is not None
+    assert nearest.entity == entities[2]
+
+
+def test_map_tree_query():
+    entities = get_test_entities()
+
+    map = Map(entities=entities)
+    tree = map.build_tree(f=entity.FlagFilter(is_collider=True))
+    test_shape = shape.Rectangle(3.0, 1.0).to_shapely(Transform2D.identity())
+    query = tree.query(test_shape)
+
+    assert len(query) == 1
+    assert query[0].entity == entities[2]
+
+
+def test_map_tree_query_nearest():
+    entities = get_test_entities()
+
+    map = Map(entities=entities)
+    tree = map.build_tree(f=entity.FlagFilter(is_collider=True))
+    test_shape = shape.Rectangle(1.0, 1.0).to_shapely(Transform2D.identity())
+    query = tree.query_nearest(test_shape)
+
+    assert len(query) == 1
+    assert query[0][0].entity == entities[2]
+    assert query[0][1] == 1.0