Skip to content

Commit

Permalink
Minor refactoring.
Browse files Browse the repository at this point in the history
  • Loading branch information
redur committed Apr 11, 2024
1 parent 6f3aee3 commit e624df5
Show file tree
Hide file tree
Showing 3 changed files with 12 additions and 10 deletions.
5 changes: 1 addition & 4 deletions src/stratigraphy/benchmark/ground_truth.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
7 changes: 3 additions & 4 deletions src/stratigraphy/util/draw.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
Expand Down
10 changes: 8 additions & 2 deletions src/stratigraphy/util/predictions.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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 = (
Expand Down

0 comments on commit e624df5

Please sign in to comment.