From 98d8a51151676c8343374cae0bc555f4967cfebb Mon Sep 17 00:00:00 2001
From: Mike Overkamp <57257130+MikeOverkamp-diag@users.noreply.github.com>
Date: Thu, 7 Jul 2022 11:17:24 +0200
Subject: [PATCH] Add segments property for overlays (#88)

Co-authored-by: James Meakin <12661555+jmsmkn@users.noreply.github.com>
---
 HISTORY.md           |  5 ++++
 panimg/models.py     | 30 +++++++++++++++++++
 pyproject.toml       |  2 +-
 tests/test_models.py | 68 ++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 104 insertions(+), 1 deletion(-)

diff --git a/HISTORY.md b/HISTORY.md
index ed5ccb7..d2800bc 100644
--- a/HISTORY.md
+++ b/HISTORY.md
@@ -1,5 +1,10 @@
 # History
 
+## 0.9.0 (2022-07-07)
+
+* Add `segments` property to the `PanImg` model containing the unique values in the image as a tuple of `int`s.
+  These are only calculated for `int` or `uint` type `SimpleITKImage`s, for any other output type `segments` are set to `None`.
+
 ## 0.8.3 (2022-06-22)
 
 * Fix installation on Windows
diff --git a/panimg/models.py b/panimg/models.py
index 32a54ce..2024b68 100644
--- a/panimg/models.py
+++ b/panimg/models.py
@@ -7,6 +7,8 @@
 from typing import Any, Dict, List, NamedTuple, Optional, Set, Tuple
 from uuid import UUID, uuid4
 
+import numpy as np
+import SimpleITK
 from pydantic import BaseModel, validator
 from pydantic.dataclasses import dataclass
 from SimpleITK import GetArrayViewFromImage, Image, WriteImage
@@ -15,6 +17,17 @@
 
 logger = logging.getLogger(__name__)
 
+MASK_TYPE_PIXEL_IDS = [
+    SimpleITK.sitkInt8,
+    SimpleITK.sitkInt16,
+    SimpleITK.sitkInt32,
+    SimpleITK.sitkInt64,
+    SimpleITK.sitkUInt8,
+    SimpleITK.sitkUInt16,
+    SimpleITK.sitkUInt32,
+    SimpleITK.sitkUInt64,
+]
+
 
 class ColorSpace(str, Enum):
     GRAY = "GRAY"
@@ -62,6 +75,8 @@ class PatientSex(str, Enum):
     "DA": lambda v: datetime.date(int(v[:4]), int(v[4:6]), int(v[6:8]))
 }
 
+MAXIMUM_SEGMENTS_LENGTH = 32
+
 
 class ExtraMetaData(NamedTuple):
     keyword: str  # DICOM tag keyword (eg. 'PatientID')
@@ -140,6 +155,7 @@ class PanImg:
     series_instance_uid: str = ""
     study_description: str = ""
     series_description: str = ""
+    segments: Optional[Tuple[int, ...]] = None
 
 
 @dataclass(frozen=True)
@@ -257,6 +273,18 @@ def add_value_range_meta_data(cls, image: Image):  # noqa: B902, N805
 
         return image
 
+    @property
+    def segments(self) -> Optional[Tuple[int, ...]]:
+        if self.image.GetPixelIDValue() not in MASK_TYPE_PIXEL_IDS:
+            return None
+
+        segments = np.unique(GetArrayViewFromImage(self.image))
+
+        if len(segments) <= MAXIMUM_SEGMENTS_LENGTH:
+            return tuple(segments)
+        else:
+            return None
+
     @property
     def color_space(self) -> ColorSpace:
         return ITK_COLOR_SPACE_MAP[self.image.GetNumberOfComponentsPerPixel()]
@@ -330,6 +358,7 @@ def save(self, output_directory: Path) -> Tuple[PanImg, Set[PanImgFile]]:
             voxel_height_mm=self.voxel_height_mm,
             voxel_depth_mm=self.voxel_depth_mm,
             eye_choice=self.eye_choice,
+            segments=self.segments,
             **self.generate_extra_metadata(),
         )
 
@@ -387,6 +416,7 @@ def save(self, output_directory: Path) -> Tuple[PanImg, Set[PanImgFile]]:
             window_center=None,
             window_width=None,
             eye_choice=self.eye_choice,
+            segments=None,
             **{md.field_name: md.default_value for md in EXTRA_METADATA},
         )
 
diff --git a/pyproject.toml b/pyproject.toml
index 9314595..06e8442 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "panimg"
-version = "0.8.3"
+version = "0.9.0"
 description = "Conversion of medical images to MHA and TIFF."
 license = "Apache-2.0"
 authors = ["James Meakin <panimg@jmsmkn.com>"]
diff --git a/tests/test_models.py b/tests/test_models.py
index 476e70f..0778069 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -1,10 +1,13 @@
 import logging
+from pathlib import Path
 
 import pytest
 
+from panimg import image_builders
 from panimg.exceptions import ValidationError
 from panimg.image_builders.metaio_utils import load_sitk_image
 from panimg.models import EXTRA_METADATA, ExtraMetaData, SimpleITKImage
+from panimg.panimg import _build_files
 from tests import RESOURCE_PATH
 
 
@@ -157,3 +160,68 @@ def test_sitk_image_value_range(
             assert not result.image.HasMetaDataKey(tag)
         else:
             assert float(result.image.GetMetaData(tag)) == pytest.approx(value)
+
+
+@pytest.mark.parametrize(
+    "src_image,builder,segments",
+    [
+        (
+            "image_min10_max10.mha",
+            image_builders.image_builder_mhd,
+            (
+                -10,
+                -9,
+                -8,
+                -7,
+                -6,
+                -5,
+                -4,
+                -3,
+                -2,
+                -1,
+                0,
+                1,
+                2,
+                3,
+                4,
+                5,
+                6,
+                7,
+                8,
+                9,
+                10,
+            ),
+        ),
+        (  # Too many values
+            "dicom_2d/cxr.dcm",
+            image_builders.image_builder_dicom,
+            None,
+        ),
+        (  # Image type is vector of ints
+            "test_rgb.png",
+            image_builders.image_builder_fallback,
+            None,
+        ),
+        (  # Tiffs are always None
+            "valid_tiff.tif",
+            image_builders.image_builder_tiff,
+            None,
+        ),
+    ],
+)
+def test_segments(
+    src_image,
+    builder,
+    segments,
+    tmpdir_factory,
+):
+    files = {RESOURCE_PATH / src_image}
+    output_dir = Path(tmpdir_factory.mktemp("output"))
+    result = _build_files(
+        builder=builder, files=files, output_directory=output_dir
+    )
+    assert result.consumed_files == files
+    assert len(result.new_images) == 1
+
+    image = result.new_images.pop()
+    assert image.segments == segments