[DAR-2228][External] - Fix for NifTI exports containing no polygons

darwin/exporter/formats/nifti.py

@@ -83,9 +83,10 @@ def export(annotation_files: Iterable[dt.AnnotationFile], output_dir: Path) -> N
             for ann in video_annotation.annotations
             if ann.annotation_class.annotation_type == "polygon"
         ]
-        populate_output_volumes_from_polygons(
-            polygon_annotations, slot_map, output_volumes
-        )
+        if polygon_annotations:
+            populate_output_volumes_from_polygons(
+                polygon_annotations, slot_map, output_volumes
+            )
         write_output_volume_to_disk(
             output_volumes, image_id=image_id, output_dir=output_dir
         )
@@ -292,25 +293,19 @@ def populate_output_volumes_from_polygons(
     annotations: List[Union[dt.Annotation, dt.VideoAnnotation]],
     slot_map: Dict,
     output_volumes: Dict,
-) -> Dict:
+):
     """
-    Exports the given ``AnnotationFile``\\s into nifti format inside of the given
-    ``output_dir``. Deletes everything within ``output_dir/masks`` before writting to it.
+    Populates the output volumes with the given polygon annotations. The annotations are converted into masks
+    and added to the corresponding volume based on the series instance UID.
 
     Parameters
     ----------
-    annotation : Union[dt.Annotation, dt.VideoAnnotation]
-        The Union of these two files used to populate the volume with
+    annotations : List[Union[dt.Annotation, dt.VideoAnnotation]]
+        List of polygon annotations used to populate the volume with
     slot_map : Dict
         Dictionary of the different slots within the annotation file
     output_volumes : Dict
-        volumes created from the build_output_volumes file
-
-    Returns
-    -------
-    volume : dict
-        Returns dict of volumes with class names as keys and volumes as values
-
+        Volumes created from the build_output_volumes file
     """
     for annotation in annotations:
         slot_name = annotation.slot_names[0]
@@ -350,7 +345,6 @@ def populate_output_volumes_from_polygons(
                 plane,
                 frame_idx,
             )
-    return volume
 
 
 def populate_output_volumes_from_raster_layer(

tests/darwin/exporter/formats/export_nifti_test.py

@@ -1,5 +1,6 @@
 import tempfile
 from pathlib import Path
+from unittest.mock import patch
 from zipfile import ZipFile
 
 import nibabel as nib
@@ -77,11 +78,55 @@ def test_video_annotation_nifti_export_mpr(team_slug_darwin_json_v2: str):
             video_annotations = list(
                 darwin_to_dt_gen(video_annotation_filepaths, False)
             )
-            nifti.export(video_annotations, output_dir=tmpdir)
+            nifti.export(video_annotations, output_dir=Path(tmpdir))
             export_im = nib.load(
                 annotations_dir / "hippocampus_001_mpr_1_test_hippo.nii.gz"
             ).get_fdata()
             expected_im = nib.load(
                 annotations_dir / "hippocampus_001_mpr_1_test_hippo.nii.gz"
             ).get_fdata()
             assert np.allclose(export_im, expected_im)
+
+
+def test_export_calls_populate_output_volumes_from_polygons(
+    team_slug_darwin_json_v2: str,
+):
+    with patch(
+        "darwin.exporter.formats.nifti.populate_output_volumes_from_polygons"
+    ) as mock:
+        with tempfile.TemporaryDirectory() as tmpdir:
+            with ZipFile("tests/data.zip") as zfile:
+                zfile.extractall(tmpdir)
+            annotations_dir = (
+                Path(tmpdir)
+                / team_slug_darwin_json_v2
+                / "nifti/releases/latest/annotations"
+            )
+            video_annotation_filepaths = [annotations_dir / "polygon_no_mask.json"]
+            video_annotations = list(
+                darwin_to_dt_gen(video_annotation_filepaths, False)
+            )
+            nifti.export(video_annotations, output_dir=Path(tmpdir))
+            mock.assert_called()
+
+
+def test_export_calls_populate_output_volumes_from_raster_layer(
+    team_slug_darwin_json_v2: str,
+):
+    with patch(
+        "darwin.exporter.formats.nifti.populate_output_volumes_from_raster_layer"
+    ) as mock:
+        with tempfile.TemporaryDirectory() as tmpdir:
+            with ZipFile("tests/data.zip") as zfile:
+                zfile.extractall(tmpdir)
+            annotations_dir = (
+                Path(tmpdir)
+                / team_slug_darwin_json_v2
+                / "nifti/releases/latest/annotations"
+            )
+            video_annotation_filepaths = [annotations_dir / "mask_no_polygon.json"]
+            video_annotations = list(
+                darwin_to_dt_gen(video_annotation_filepaths, False)
+            )
+            nifti.export(video_annotations, output_dir=Path(tmpdir))
+            mock.assert_called()