Merge pull request #42 from MannLabs/shape-attributes-1

Custom shapes attributes

Author: sophiamaedler

src/lmd/lib.py

@@ -31,7 +31,7 @@
 from skimage.segmentation import find_boundaries
 
 from pathlib import Path
-from typing import Callable, Optional, Union, Iterable
+from typing import Callable, Optional, Union, Iterable, Any
 
 import geopandas as gpd
 import matplotlib.pyplot as plt
@@ -263,7 +263,7 @@ def add_shape(self, shape: Shape):
             TypeError("Provided shape is not of type Shape")
 
     def new_shape(
-        self, points: np.ndarray, well: Optional[str] = None, name: Optional[str] = None
+        self, points: np.ndarray, well: Optional[str] = None, name: Optional[str] = None, **custom_attributes
     ):
         """Directly create a new Shape in the current collection.
 
@@ -273,13 +273,17 @@ def new_shape(
             well: Well in which to sort the shape after cutting. For example A1, A2 or B3.
 
             name: Name of the shape.
-        """
 
+            custom_attributes: Custom shape metadata that can be added as additional xml-element to the shape.
+                All values are converted to strings.
+
+        """
         to_add = Shape(
             points,
             well=well,
             name=name,
             orientation_transform=self.orientation_transform,
+            **custom_attributes
         )
         self.add_shape(to_add)
 
@@ -319,22 +323,29 @@ def to_geopandas(self, *attrs: str) -> gpd.GeoDataFrame:
         .. code-block:: python
             # Generate collection
             collection = pylmd.Collection()
-            shape = pylmd.Shape(np.array([[ 0,  0], [ 0, -1], [ 1,  0], [ 0,  0]]), well="A1", name="Shape_1", orientation_transform=None)
+            shape = pylmd.Shape(
+                    np.array([[ 0,  0], [ 0, -1], [ 1,  0], [ 0,  0]]), 
+                    well="A1", 
+                    name="Shape_1", 
+                    metadata1="A",
+                    metadata2="B",
+                    orientation_transform=None
+                )
             collection.add_shape(shape)
 
             # Get geopandas object
             collection.to_geopandas()
             >       geometry
                 0   POLYGON ((0 0, 0 -1, 1 0, 0 0))
 
-            collection.to_geopandas("well", "name")
-            >   well     name                         geometry
-                0   A1  Shape_1  POLYGON ((0 0, 0 -1, 1 0, 0 0))
+            collection.to_geopandas("well", "name", "metadata1", "metadata2")
+            >       well    name            metadata1 metadata2  geometry
+                0   A1      Shape_1         A         B          POLYGON ((0 0, 0 -1, 1 0, 0 0))
         """
         metadata = (
             pd.DataFrame(
                 [
-                    [shape.__getattribute__(att) for att in attrs]
+                    [shape.get_shape_annotation(att) for att in attrs]
                     for shape in self.shapes
                 ],
                 columns=attrs,
@@ -393,6 +404,7 @@ def load_geopandas(
             well_column: Optional[str] = None,
             calibration_points: Optional[np.ndarray] = None, 
             global_coordinates: Optional[int] = None,
+            custom_attribute_columns: str | list[str] | None = None,
         ) -> None:
         """Create collection from a geopandas dataframe
         
@@ -402,6 +414,8 @@ def load_geopandas(
             well_column (str, optional): Column storing of well id as additional metadata
             calibration_points (np.ndarray, optional): Calibration points of collection 
             global_coordinates (int, optional): Number of global coordinates
+            custom_attribute_columns Custom shape metadata that will be added as additional xml-element to the shape. 
+                Can be column name, list of column names or None
 
         Example:
 
@@ -433,11 +447,17 @@ def load_geopandas(
         if global_coordinates is not None:
             self.global_coordinates = global_coordinates
 
+        if custom_attribute_columns is None:
+            custom_attribute_columns = []
+        if isinstance(custom_attribute_columns, str):
+            custom_attribute_columns = [custom_attribute_columns]
+
         self.shapes = [
             Shape(
                 points=np.array(row[geometry_column].exterior.coords), 
                 name=row[name_column] if name_column is not None else None,
                 well=row[well_column] if well_column is not None else None,
+                **{att: row[att] for att in custom_attribute_columns}
             )
             for _, row in gdf.iterrows()
         ]
@@ -542,6 +562,7 @@ def __init__(
         well: Optional[str] = None,
         name: Optional[str] = None,
         orientation_transform=None,
+        **custom_attributes: dict[str, str]
     ):
         """Class for creating a single shape.
 
@@ -551,8 +572,10 @@ def __init__(
             well: Well in which to sort the shape after cutting. For example A1, A2 or B3.
 
             name: Name of the shape.
-        """
 
+            custom_attributes: Custom shape metadata that will be added as additional xml-element to the shape
+                Values be implicitly converted to strings.
+        """
         # Orientation transform of shapes
         self.orientation_transform: Optional[np.ndarray] = orientation_transform
 
@@ -569,6 +592,8 @@ def __init__(
         self.name: Optional[str] = name
         self.well: Optional[str] = well
 
+        self.custom_attributes = custom_attributes
+
     def from_xml(self, root):
         """Load a shape from an XML shape node. Used internally for reading LMD generated XML files.
         
@@ -599,10 +624,14 @@ def from_xml(self, root):
                 points[point_id, 1] = int(child.text)  
             elif child.tag == "CapID":
                 self.well = str(child.text)
+            else: 
+                if child.tag in self.custom_attributes:
+                    warnings.warn(f"Shape attribute {child.tag} already found in shape, overwrite", stacklevel=1)
+                self.custom_attributes[child.tag] = child.text
 
         self.points = np.array(points)
 
-    def to_xml(self, id: int, orientation_transform: np.ndarray, scale: int):
+    def to_xml(self, id: int, orientation_transform: np.ndarray, scale: int, *, write_custom_attributes: bool = True):
         """Generate XML shape node needed internally for export.
 
         Args:
@@ -612,6 +641,8 @@ def to_xml(self, id: int, orientation_transform: np.ndarray, scale: int):
 
             scale (int): Scalling factor used to enable higher decimal precision.
 
+            write_custom_attributes: Write custom attributes to xml file
+
         Note:
             If the Shape has a custom orientation_transform defined, the custom orientation_transform is applied at this point. If not, the oritenation_transform passed by the parent Collection is used. This highlights an important difference between the Shape and Collection class. The Collection will always has an orientation transform defined and will use `np.eye(2)` by default. The Shape object can have a orientation_transform but can also be set to `None` to use the Collection value.
 
@@ -633,6 +664,12 @@ def to_xml(self, id: int, orientation_transform: np.ndarray, scale: int):
             cap_id = ET.SubElement(shape, "CapID")
             cap_id.text = self.well
 
+        if write_custom_attributes:
+            for attribute_name, attribute_value in self.custom_attributes.items():
+                custom_attribute = ET.SubElement(shape, attribute_name)
+                # xml only accepts string values
+                custom_attribute.text = str(attribute_value)
+
         # write points
         for i, point in enumerate(transformed_points):
             id = i + 1
@@ -643,6 +680,28 @@ def to_xml(self, id: int, orientation_transform: np.ndarray, scale: int):
             y.text = "{}".format(np.floor(point[1]).astype(int))
 
         return shape
+    
+    def get_shape_annotation(self, name: str) -> Any | None:
+        """Retrieve the value of an attribute from either instance attributes
+         or custom attributes by name.
+
+         Searches for the attribute by name in the 1) instance attributes
+         2) custom attributes, or 3) issues a warning and returns None
+
+        Args:
+            name (str): The name of the attribute to retrieve.
+
+        Returns:
+            Any | None: The value of the attribute if found, otherwise None.
+        """
+        if name in self.__dict__:
+            return getattr(self, name)
+        elif name in self.custom_attributes:
+            return self.custom_attributes.get(name)
+        else:
+            warnings.warn(f"Attribute {name} not found in shape attributes. Returning None.")
+            return None
+
 
     def to_shapely(self):
         return shapely.Polygon(self.points)

src/lmd/lmd_test.py

@@ -8,22 +8,29 @@
 import geopandas as gpd 
 import shapely 
 from lxml import etree as ET
+import pytest
 
 def test_collection():
     calibration = np.array([[0, 0], [0, 100], [50, 50]])
     my_first_collection = Collection(calibration_points = calibration)
 
+
 def test_shape():
     rectangle_coordinates = np.array([[10,10], [40,10], [40,40], [10,40], [10,10]])
     rectangle = Shape(rectangle_coordinates)
 
+
 def test_shape_from_xml():
+    """Read a minimal xml representation of a cell shape and associated metadata"""
     # Define shape in xml
     shape_xml = """
     <Shape_1>
         <PointCount>3</PointCount>
         <CapID>A1</CapID>
         <TEST>this is a test</TEST>
+        <test2>1</test2>
+        <test3>3.1415</test3>
+
         <X_1>0</X_1>
         <Y_1>-0</Y_1>
         <X_2>0</X_2>
@@ -41,6 +48,10 @@ def test_shape_from_xml():
     shape.from_xml(shape_xml)
     assert (shape.points == np.array([[ 0,  0], [ 0, -1], [ 1,  0]])).all()
     assert shape.well == "A1"
+    assert shape.custom_attributes["TEST"] == "this is a test"
+    assert shape.custom_attributes["test2"] == "1"
+    assert shape.custom_attributes["test3"] == "3.1415"
+
 
 def test_plotting():
     calibration = np.array([[0, 0], [0, 100], [50, 50]])
@@ -60,31 +71,68 @@ def test_plotting():
 
     my_first_collection.plot(calibration = True)
 
-def test_collection_load_geopandas():
-    gdf = gpd.GeoDataFrame(
-        data={"well": ["A1"], "name": "my_shape"},
+
+@pytest.fixture
+def geopandas_collection():
+    """Geopandas shape collection with both controlled (name, well) and custom metadata"""
+    return gpd.GeoDataFrame(
+        data={"well": ["A1"], "name": "my_shape", "string_attribute": "a"},
         geometry=[shapely.Polygon([[0, 0], [0, 1], [1, 0], [0, 0]])]
     )
 
+@pytest.mark.parametrize(
+        ("well_column", "name_column", "custom_attributes"),
+        [
+            ("well", None, None),
+            (None, "well", None),
+            (None, None, "string_attribute"),
+            ("well", "name", None),
+            ("well", "name", "string_attribute"),
+
+        ]
+)
+def test_collection_load_geopandas(
+    geopandas_collection: gpd.GeoDataFrame, 
+    well_column: str, 
+    name_column: str,
+    custom_attributes: list[str]
+    ) -> None:
 
     # Export well metadata
     c = Collection(calibration_points=np.array([[-1, -1], [1, 1], [0, 1]]))
     calibration_points_old = c.calibration_points
-    c.load_geopandas(gdf, well_column="well", name_column="name")
-    assert c.to_geopandas("well", "name").equals(gdf)
+    c.load_geopandas(
+        geopandas_collection, 
+        well_column=well_column, 
+        name_column=name_column, 
+        custom_attribute_columns=custom_attributes
+    )
+    
+    all_columns = [
+        col for col in (well_column, custom_attributes) if col is not None
+    ]
+
+    assert c.to_geopandas(*all_columns).equals(geopandas_collection[[*all_columns, "geometry"]])
     assert (c.calibration_points == calibration_points_old).all()
 
     # Overwrite calibration points 
     c = Collection(calibration_points=np.array([[-1, -1], [1, 1], [0, 1]]))
     calibration_points_new = np.array([[0, 0], [100, 0], [0, 100]])
-    c.load_geopandas(gdf, well_column="well", name_column="name", calibration_points=calibration_points_new)
-    assert c.to_geopandas("well", "name").equals(gdf)
+
+    c.load_geopandas(
+        geopandas_collection, 
+        calibration_points=calibration_points_new, 
+        well_column=well_column, 
+        name_column=name_column, 
+        custom_attribute_columns=custom_attributes
+    )
+    assert c.to_geopandas(*all_columns).equals(geopandas_collection[[*all_columns, "geometry"]])
     assert (c.calibration_points == calibration_points_new).all()
 
     # Do not export well metadata
     c = Collection(calibration_points=np.array([[-1, -1], [1, 1], [0, 1]]))
-    c.load_geopandas(gdf)
-    assert c.to_geopandas().equals(gdf.drop(columns=["well", "name"]))
+    c.load_geopandas(geopandas_collection)
+    assert c.to_geopandas().equals(geopandas_collection[["geometry"]])
 
 def test_collection_save():
     calibration = np.array([[0, 0], [0, 100], [50, 50]])
@@ -98,7 +146,8 @@ def test_collection_save():
     my_first_collection.add_shape(rectangle)
 
     my_first_collection.save("first_collection.xml")
-    
+
+
 def test_tools_square():
     calibration = np.array([[0, 0], [0, 100], [50, 50]])
     my_first_collection = Collection(calibration_points = calibration)
@@ -108,7 +157,8 @@ def test_tools_square():
 
     my_square = tools.rectangle(10, 10, offset=(30,30))
     my_first_collection.add_shape(my_square)
-                                  
+
+                            
 def test_glyphs():
     calibration = np.array([[0, 0], [0, 100], [50, 50]])
     my_first_collection = Collection(calibration_points = calibration)
@@ -135,6 +185,7 @@ def test_text():
     identifier_3 = tools.text('0123456789-_ABCDEFGHI', offset=np.array([60, 40]), rotation = -np.pi/4)
     my_first_collection.join(identifier_3)                 
 
+
 def test_segmentation_loader():
 
     package_base_path = pathlib.Path(__file__).parent.parent.parent.resolve().absolute()