diff --git a/darwin/exporter/formats/nifti.py b/darwin/exporter/formats/nifti.py index 64dcb353b..65576952a 100644 --- a/darwin/exporter/formats/nifti.py +++ b/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( diff --git a/tests/darwin/exporter/formats/export_nifti_test.py b/tests/darwin/exporter/formats/export_nifti_test.py index d21c8d596..648e1ed72 100644 --- a/tests/darwin/exporter/formats/export_nifti_test.py +++ b/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,7 +78,7 @@ 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() @@ -85,3 +86,47 @@ def test_video_annotation_nifti_export_mpr(team_slug_darwin_json_v2: str): 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() diff --git a/tests/data.zip b/tests/data.zip index 5713fff0f..dedd46487 100644 Binary files a/tests/data.zip and b/tests/data.zip differ