diff --git a/src/stratigraphy/benchmark/ground_truth.py b/src/stratigraphy/benchmark/ground_truth.py
index 5673d237..c98904cb 100644
--- a/src/stratigraphy/benchmark/ground_truth.py
+++ b/src/stratigraphy/benchmark/ground_truth.py
@@ -11,10 +11,7 @@
 
 
 class GroundTruthForFile:
-    """Ground truth data for a single file.
-
-    Evaluates the prediction against the ground truth.
-    """
+    """Ground truth data for a single file."""
 
     def __init__(self, ground_truth_layers: list):
         self.layers = ground_truth_layers
diff --git a/src/stratigraphy/util/draw.py b/src/stratigraphy/util/draw.py
index 56228768..7bc07702 100644
--- a/src/stratigraphy/util/draw.py
+++ b/src/stratigraphy/util/draw.py
@@ -72,17 +72,16 @@ def draw_material_descriptions(page: fitz.Page, layers: LayerPrediction) -> None
         layers (LayerPrediction): The predictions for the page.
     """
     for index, layer in enumerate(layers):
-        material_description_layer = layer.material_description
-        if material_description_layer.rect is not None:
+        if layer.material_description.rect is not None:
             fitz.utils.draw_rect(
                 page,
-                fitz.Rect(material_description_layer.rect) * page.derotation_matrix,
+                fitz.Rect(layer.material_description.rect) * page.derotation_matrix,
                 color=fitz.utils.getColor("orange"),
             )
         draw_layer(
             page=page,
             interval=layer.depth_interval,  # None if no depth interval
-            layer=material_description_layer,
+            layer=layer.material_description,
             index=index,
             is_correct=layer.material_is_correct,  # None if no ground truth
         )
diff --git a/src/stratigraphy/util/predictions.py b/src/stratigraphy/util/predictions.py
index 0052b61b..2f1a311b 100644
--- a/src/stratigraphy/util/predictions.py
+++ b/src/stratigraphy/util/predictions.py
@@ -10,7 +10,10 @@
 
 
 class MaterialDescriptionPrediction:
-    """A class to represent a material description prediction."""
+    """A class to represent a material description prediction.
+
+    TODO: Check if this class can be replaced.
+    """
 
     def __init__(self, text: str, rect: list, lines: list):
         self.text = text
@@ -19,7 +22,10 @@ def __init__(self, text: str, rect: list, lines: list):
 
 
 class DepthIntervalPrediction:
-    """A class to represent a depth interval prediction."""
+    """A class to represent a depth interval prediction.
+
+    TODO: Could be replaced by BoundaryInterval from interval.py. Then we could use line_anchor and background_rect.
+    """
 
     def __init__(self, start: dict, end: dict):
         self.start = (