Merge pull request #265 from uncharted-lara/fix-for-missing-segmentat…

…ions fix: handle cases where input data is missing `SEGMENTATION_OUTPUT_KEY`
DARPA-CRITICALMAAS · Aug 15, 2024 · b975870 · b975870
2 parents 7880a35 + dfcea7b
commit b975870
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/tasks/metadata_extraction/metadata_extraction.py b/tasks/metadata_extraction/metadata_extraction.py
@@ -367,7 +367,7 @@ def run(self, input: TaskInput) -> TaskResult:
                 metadata.utm_zone = str(self._extract_utm_zone(metadata))
 
             # compute map shape from the segmentation output
-            segments = input.data[SEGMENTATION_OUTPUT_KEY]
+            segments = input.data.get(SEGMENTATION_OUTPUT_KEY, None)
             metadata.map_shape = self._compute_shape(segments)
 
             # compute map chroma from the image
@@ -779,6 +779,7 @@ def _compute_shape(self, segments) -> MapShape:
         Returns:
             MapShape: The shape of the map
         """
+        map_shape = MapShape.UNKNOWN
         if segments:
             map_segmentation = MapSegmentation.model_validate(segments)
             for segment in map_segmentation.segments:

diff --git a/tasks/metadata_extraction/text_filter.py b/tasks/metadata_extraction/text_filter.py
@@ -47,7 +47,9 @@ def run(self, input: TaskInput) -> TaskResult:
         doc_text = DocTextExtraction.model_validate(text_data)
 
         # get map segments
-        segments = input.data[SEGMENTATION_OUTPUT_KEY]
+        segments = input.data.get(
+            SEGMENTATION_OUTPUT_KEY, {"doc_id": input.raster_id, "segments": []}
+        )
         map_segmentation = MapSegmentation.model_validate(segments)
 
         output_text: Dict[str, TextExtraction] = {}