Add segments property for overlays (#88)

Co-authored-by: James Meakin <[email protected]>

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

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},
  )
 

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 <[email protected]>"]

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