From f3cb324938758a1ea6ad3cd35056f87fafc4c6d8 Mon Sep 17 00:00:00 2001 From: Andras Lasso Date: Sat, 4 May 2024 03:00:43 -0400 Subject: [PATCH] Simplify segmentation API and add support for terminologies Segmentation methods now take a single segmentation object (dict) that contains both voxels and metadata. Added functions to: - create segmentation object from scratch - write segmentation object to file - extract segments from a segmentation based on terminology codes --- README.md | 99 +++- slicerio/__init__.py | 5 +- slicerio/segmentation.py | 786 +++++++++++++++++++++++++--- slicerio/tests/test_segmentation.py | 130 ++++- 4 files changed, 918 insertions(+), 102 deletions(-) diff --git a/README.md b/README.md index 6847257..969c1b1 100644 --- a/README.md +++ b/README.md @@ -24,20 +24,49 @@ pip install slicerio import slicerio import json -segmentation_info = slicerio.read_segmentation_info("Segmentation.seg.nrrd") +segmentation = slicerio.read_segmentation("path/to/Segmentation.seg.nrrd", skip_voxels=True) -number_of_segments = len(segmentation_info["segments"]) +number_of_segments = len(segmentation["segments"]) print(f"Number of segments: {number_of_segments}") -segment_names = slicerio.segment_names(segmentation_info) +segment_names = slicerio.segment_names(segmentation) print(f"Segment names: {', '.join(segment_names)}") -segment0 = slicerio.segment_from_name(segmentation_info, segment_names[0]) +segment0 = slicerio.segment_from_name(segmentation, segment_names[0]) print("First segment info:\n" + json.dumps(segment0, sort_keys=False, indent=4)) ``` ### Extract selected segments with chosen label values +#### Extract segments by terminology + +Example for getting a 3D NRRD file that has label values assigned based on standard terminology codes. +Terminology is a `dict` that must specify `category` and `type` codes and may optionally also specify `typeModifier`, `anatomicRegion`, and `anatomicRegionModifier`. Each code is specifed by a triplet of "coding scheme designator", "code value", "code meaning" in a list. + +Coding scheme designator is typically `SCT` (SNOMED-CT) for clinical images. You can find codes in the [SNOMED-CT browser](https://browser.ihtsdotools.org/). When code exists for "entire X" and "structure of X" then always use the "structure" code ("entire" code has a very strict meaning that is rarely applicable in practice). + +Code meaning (third component of codes, such as "Anatomical Structure", "Ribs", "Right") is informational only, it can be used for troubleshooting or displayed to the user, but it is ignored in information processing (e.g., two codes match if their coding scheme designator and code value are the same even if code meaning is different). + +```python +import slicerio +import nrrd + +input_filename = "path/to/Segmentation.seg.nrrd" +output_filename = "path/to/SegmentationExtracted.seg.nrrd" +segments_to_labels = [ + ({"category": ["SCT", "123037004", "Anatomical Structure"], "type": ["SCT", "113197003", "Ribs"]}, 1), + ({"category": ["SCT", "123037004", "Anatomical Structure"], "type": ["SCT", "39607008", "Lung"], "typeModifier": ["SCT", "24028007", "Right"]}, 3) + ] + +segmentation = slicerio.read_segmentation(input_filename) +extracted_segmentation = slicerio.extract_segments(segmentation, segments_to_labels) +slicerio.write_segmentation(extracted_segmentation, output_filename) +``` + +#### Extract segments by name + +It is strongly recommended to look up segments by standard terminology codes instead of segment name, as spelling errors and inconsistent use of names often causes mismatch. + ```python import slicerio import nrrd @@ -46,9 +75,65 @@ input_filename = "path/to/Segmentation.seg.nrrd" output_filename = "path/to/SegmentationExtracted.seg.nrrd" segment_names_to_labels = [("ribs", 10), ("right lung", 12), ("left lung", 6)] -voxels, header = nrrd.read(input_filename) -extracted_voxels, extracted_header = slicerio.extract_segments(voxels, header, segmentation_info, segment_names_to_labels) -nrrd.write(output_filename, extracted_voxels, extracted_header) +segmentation = slicerio.read_segmentation(input_filename) +extracted_segmentation = slicerio.extract_segments(segmentation, segment_names_to_labels) +slicerio.write_segmentation(extracted_segmentation, output_filename) +``` + +### Create segmentation file from numpy array + +```python +# Create segmentation with two labels (1, 3) +voxels = np.zeros([100, 120, 150]) +voxels[30:50, 20:60, 70:100] = 1 +voxels[70:90, 80:110, 60:110] = 3 + +# Image geometry +spacing = [0.5, 0.5, 0.8] +origin = [10, 30, 15] + +segmentation = { + "voxels": voxels, + "image": { + "encoding": "gzip", + "ijkToLPS": [[ spacing[0], 0., 0., origin[0]], + [ 0., spacing[1], 0., origin[1]], + [ 0., 0., spacing[2], origin[2]], + [ 0., 0., 0., 1. ]] + }, + "segmentation": { + "containedRepresentationNames": ["Binary labelmap", "Closed surface"], + # "masterRepresentation": "Binary labelmap", + # "referenceImageExtentOffset": [0, 0, 0], + }, + "segments": [ + { + "id": "Segment_1", + "labelValue": 1, + "layer": 0, + "color": [0.9, 0.9, 0.6], + "name": "ribs", + "terminology": { + "contextName": "Segmentation category and type - 3D Slicer General Anatomy list", + "category": ["SCT", "123037004", "Anatomical Structure"], + "type": ["SCT", "113197003", "Rib"] } + }, + { + "id": "Segment_2", + "labelValue": 3, + "layer": 0, + "color": [0.9, 0.9, 0.6], + "name": "spine", + "status": "inprogress", + "terminology": { + "contextName": "Segmentation category and type - 3D Slicer General Anatomy list", + "category": ["SCT", "123037004", "Anatomical Structure"], + "type": ["SCT", "122494005", "Cervical spine"] } + }, + ] +} + +slicerio.write_segmentation(segmentation, "path/to/Segmentation.seg.nrrd") ``` ### View files in 3D Slicer diff --git a/slicerio/__init__.py b/slicerio/__init__.py index 8019cb3..c35f0c9 100644 --- a/slicerio/__init__.py +++ b/slicerio/__init__.py @@ -12,14 +12,15 @@ """ -from .segmentation import extract_segments, read_segmentation_info, segment_from_name, segment_names +from .segmentation import extract_segments, read_segmentation, write_segmentation, segment_from_name, segment_names from .data_helper import get_testdata_file from ._version import __version__, __version_info__ __all__ = [ 'extract_segments', 'get_testdata_file', - 'read_segmentation_info', + 'read_segmentation', + 'write_segmentation', 'segment_from_name', 'segment_names', '__version__', diff --git a/slicerio/segmentation.py b/slicerio/segmentation.py index eed8f4c..a34c9bf 100644 --- a/slicerio/segmentation.py +++ b/slicerio/segmentation.py @@ -1,7 +1,34 @@ # -*- coding: utf-8 -*- -def read_terminology_entry(terminology_entry): - terminology_items = terminology_entry.split('~') +def terminology_entry_from_string(terminology_str): + """Converts a terminology string to a dict. + + Example terminology string: + + Segmentation category and type - 3D Slicer General Anatomy list + ~SCT^49755003^Morphologically Altered Structure + ~SCT^4147007^Mass + ~^^ + ~Anatomic codes - DICOM master list + ~SCT^23451007^Adrenal gland + ~SCT^24028007^Right + + Resulting dict: + + { + 'contextName': 'Segmentation category and type - 3D Slicer General Anatomy list', + 'category': ['SCT', '49755003', 'Morphologically Altered Structure'], + 'type': ['SCT', '4147007', 'Mass'], + 'anatomicContextName': 'Anatomic codes - DICOM master list', + 'anatomicRegion': ['SCT', '23451007', 'Adrenal gland'], + 'anatomicRegionModifier': ['SCT', '24028007', 'Right'] + } + + Specification of terminology entry string is available at + https://slicer.readthedocs.io/en/latest/developer_guide/modules/segmentations.html#terminologyentry-tag + """ + + terminology_items = terminology_str.split('~') terminology = {} terminology['contextName'] = terminology_items[0] @@ -12,9 +39,11 @@ def read_terminology_entry(terminology_entry): if any(item != '' for item in typeModifier): terminology['typeModifier'] = typeModifier + anatomicContextName = terminology_items[4] + if anatomicContextName: + terminology['anatomicContextName'] = anatomicContextName anatomicRegion = terminology_items[5].split("^") if any(item != '' for item in anatomicRegion): - terminology['anatomicContextName'] = terminology_items[4] terminology['anatomicRegion'] = anatomicRegion anatomicRegionModifier = terminology_items[6].split("^") if any(item != '' for item in anatomicRegionModifier): @@ -22,113 +51,706 @@ def read_terminology_entry(terminology_entry): return terminology +def terminology_entry_to_string(terminology): + """Converts a terminology dict to string. + """ + terminology_str = "" -def read_segmentation_info(filename): - import nrrd - header = nrrd.read_header(filename) - segmentation_info = {} - segments = [] - segment_index = 0 + if 'contextName' in terminology: + terminology_str += terminology['contextName'] + else: + terminology_str += "" + terminology_str += '~' + '^'.join(terminology['category']) + terminology_str += '~' + '^'.join(terminology['type']) + if 'typeModifier' in terminology: + typeModifier = terminology['typeModifier'] + else: + typeModifier = ['', '', ''] + terminology_str += '~' + '^'.join(typeModifier) + + if 'anatomicContextName' in terminology: + terminology_str += "~" + terminology['anatomicContextName'] + else: + terminology_str += "~" + if 'anatomicRegion' in terminology: + anatomic_region = terminology['anatomicRegion'] + else: + anatomic_region = ['', '', ''] + terminology_str += '~' + '^'.join(anatomic_region) + if 'anatomicRegionModifier' in terminology: + anatomic_region_modifier = terminology['anatomicRegionModifier'] + else: + anatomic_region_modifier = ['', '', ''] + terminology_str += '~' + '^'.join(anatomic_region_modifier) + + return terminology_str + + +def generate_unique_segment_id(existing_segment_ids): + """Generate a unique segment ID, i.e., an ID that is not among existing_segment_ids.""" + i = 1 while True: - prefix = "Segment{0}_".format(segment_index) - if not prefix + "ID" in header.keys(): - break - segment = {} - segment["index"] = segment_index - segment["color"] = [float(i) for i in header[prefix + "Color"].split(" ")] # Segment0_Color:=0.501961 0.682353 0.501961 - segment["colorAutoGenerated"] = int(header[prefix + "ColorAutoGenerated"]) != 0 # Segment0_ColorAutoGenerated:=1 - segment["extent"] = [int(i) for i in header[prefix + "Extent"].split(" ")] # Segment0_Extent:=68 203 53 211 24 118 - segment["id"] = header[prefix + "ID"] # Segment0_ID:=Segment_1 - segment["labelValue"] = int(header[prefix + "LabelValue"]) # Segment0_LabelValue:=1 - segment["layer"] = int(header[prefix + "Layer"]) # Segment0_Layer:=0 - segment["name"] = header[prefix + "Name"] # Segment0_Name:=Segment_1 - segment["nameAutoGenerated"] = int(header[prefix + "NameAutoGenerated"]) != 0 # Segment0_NameAutoGenerated:=1 + segment_id = "Segment_" + str(i) + if segment_id not in existing_segment_ids: + return segment_id + i += 1 + + +def read_segmentation(filename, skip_voxels=False): + """Read segmentation metadata from a .seg.nrrd file and store it in a dict. + + Example header: + + NRRD0004 + # Complete NRRD file format specification at: + # http://teem.sourceforge.net/nrrd/format.html + type: unsigned char + dimension: 3 + space: left-posterior-superior + sizes: 128 128 34 + space directions: (-3.04687595367432,0,0) (0,-3.04687595367432,0) (0,0,9.9999999999999964) + kinds: domain domain domain + encoding: gzip + space origin: (193.09599304199222,216.39599609374994,-340.24999999999994) + Segment0_Color:=0.992157 0.909804 0.619608 + Segment0_ColorAutoGenerated:=1 + Segment0_Extent:=0 124 0 127 0 33 + Segment0_ID:=Segment_1 + Segment0_LabelValue:=1 + Segment0_Layer:=0 + Segment0_Name:=ribs + Segment0_NameAutoGenerated:=1 + Segment0_Tags:=Segmentation.Status:inprogress|TerminologyEntry:Segmentation category and type - 3D Slicer General Anatomy list~SCT^123037004^Anatomical Structure~SCT^113197003^Rib~^^~Anatomic codes - DICOM master list~^^~^^| + Segment1_Color:=1 1 0.811765 + Segment1_ColorAutoGenerated:=1 + Segment1_Extent:=0 124 0 127 0 33 + Segment1_ID:=Segment_2 + Segment1_LabelValue:=2 + Segment1_Layer:=0 + Segment1_Name:=cervical vertebral column + Segment1_NameAutoGenerated:=1 + Segment1_Tags:=Segmentation.Status:inprogress|TerminologyEntry:Segmentation category and type - 3D Slicer General Anatomy list~SCT^123037004^Anatomical Structure~SCT^122494005^Cervical spine~^^~Anatomic codes - DICOM master list~^^~^^| + Segment2_Color:=0.886275 0.792157 0.52549 + Segment2_ColorAutoGenerated:=1 + Segment2_Extent:=0 124 0 127 0 33 + Segment2_ID:=Segment_3 + Segment2_LabelValue:=3 + Segment2_Layer:=0 + Segment2_Name:=thoracic vertebral column + Segment2_NameAutoGenerated:=1 + Segment2_Tags:=Some field:some value|Segmentation.Status:inprogress|TerminologyEntry:Segmentation category and type - 3D Slicer General Anatomy list~SCT^123037004^Anatomical Structure~SCT^122495006^Thoracic spine~^^~Anatomic codes - DICOM master list~^^~^^| + Segmentation_ContainedRepresentationNames:=Binary labelmap|Closed surface| + Segmentation_ConversionParameters:=Collapse labelmaps|1|Merge the labelmaps into as few shared labelmaps as possible 1 = created labelmaps will be shared if possible without overwriting each other.&Compute surface normals|1|Compute surface normals. 1 (default) = surface normals are computed. 0 = surface normals are not computed (slightly faster but produces less smooth surface display).&Crop to reference image geometry|0|Crop the model to the extent of reference geometry. 0 (default) = created labelmap will contain the entire model. 1 = created labelmap extent will be within reference image extent.&Decimation factor|0.0|Desired reduction in the total number of polygons. Range: 0.0 (no decimation) to 1.0 (as much simplification as possible). Value of 0.8 typically reduces data set size by 80% without losing too much details.&Default slice thickness|0.0|Default thickness for contours if slice spacing cannot be calculated.&End capping|1|Create end cap to close surface inside contours on the top and bottom of the structure.\n0 = leave contours open on surface exterior.\n1 (default) = close surface by generating smooth end caps.\n2 = close surface by generating straight end caps.&Fractional labelmap oversampling factor|1|Determines the oversampling of the reference image geometry. All segments are oversampled with the same value (value of 1 means no oversampling).&Joint smoothing|0|Perform joint smoothing.&Oversampling factor|1|Determines the oversampling of the reference image geometry. If it's a number, then all segments are oversampled with the same value (value of 1 means no oversampling). If it has the value "A", then automatic oversampling is calculated.&Reference image geometry|3.0468759536743195;0;0;-193.0959930419922;0;3.0468759536743195;0;-216.39599609374994;0;0;9.999999999999998;-340.24999999999994;0;0;0;1;0;127;0;127;0;33;|Image geometry description string determining the geometry of the labelmap that is created in course of conversion. Can be copied from a volume, using the button.&Smoothing factor|-0.5|Smoothing factor. Range: 0.0 (no smoothing) to 1.0 (strong smoothing).&Threshold fraction|0.5|Determines the threshold that the closed surface is created at as a fractional value between 0 and 1.& + Segmentation_MasterRepresentation:=Binary labelmap + Segmentation_ReferenceImageExtentOffset:=0 0 0 + + Example header in case of overlapping segments: + + NRRD0004 + # Complete NRRD file format specification at: + # http://teem.sourceforge.net/nrrd/format.html + type: unsigned char + dimension: 4 + space: left-posterior-superior + sizes: 5 256 256 130 + space directions: none (0,1,0) (0,0,-1) (-1.2999954223632812,0,0) + kinds: list domain domain domain + encoding: gzip + space origin: (86.644897460937486,-133.92860412597656,116.78569793701172) + Segment0_... + + Returned segmentation object: + + { + "voxels": (numpy array of voxel values), + "ijkToLPS": [[ -3.04687595, 0. , 0. , 193.09599304], + [ 0. , -3.04687595, 0. , 216.39599609], + [ 0. , 0. , 10. , -340.25 ], + [ 0. , 0. , 0. , 1. ]], + "encoding": "gzip", + "containedRepresentationNames": ["Binary labelmap", "Closed surface"], + "conversionParameters": [ + {"name": "Collapse labelmaps", "value": "1", "description": "Merge the labelmaps into as few shared labelmaps as possible 1 = created labelmaps will be shared if possible without overwriting each other."}, + {"name": "Compute surface normals", "value": "1", "description": "Compute surface normals. 1 (default) = surface normals are computed. 0 = surface normals are not computed (slightly faster but produces less smooth surface display)."}, + {"name": "Crop to reference image geometry", "value": "0", "description": "Crop the model to the extent of reference geometry. 0 (default) = created labelmap will contain the entire model. 1 = created labelmap extent will be within reference image extent."}, + {"name": "Decimation factor", "value": "0.0", "description": "Desired reduction in the total number of polygons. Range: 0.0 (no decimation) to 1.0 (as much simplification as possible). Value of 0.8 typically reduces data set size by 80% without losing too much details."}, + {"name": "Default slice thickness", "value": "0.0", "description": "Default thickness for contours if slice spacing cannot be calculated."}, + {"name": "End capping", "value": "1", "description": "Create end cap to close surface inside contours on the top and bottom of the structure.\n0 = leave contours open on surface exterior.\n1 (default) = close surface by generating smooth end caps.\n2 = close surface by generating straight end caps."}, + {"name": "Fractional labelmap oversampling factor", "value": "1", "description": "Determines the oversampling of the reference image geometry. All segments are oversampled with the same value (value of 1 means no oversampling)."}, + {"name": "Joint smoothing", "value": "0", "description": "Perform joint smoothing."}, + {"name": "Oversampling factor", "value": "1", "description": "Determines the oversampling of the reference image geometry. If it's a number, then all segments are oversampled with the same value (value of 1 means no oversampling). If it has the value \"A\", then automatic oversampling is calculated."}, + {"name": "Reference image geometry", "value": "3.0468759536743195;0;0;-193.0959930419922;0;3.0468759536743195;0;-216.39599609374994;0;0;9.999999999999998;-340.24999999999994;0;0;0;1;0;127;0;127;0;33;", "description": "Image geometry description string determining the geometry of the labelmap that is created in course of conversion. Can be copied from a volume, using the button."}, + {"name": "Smoothing factor", "value": "-0.5", "description": "Smoothing factor. Range: 0.0 (no smoothing) to 1.0 (strong smoothing)."}, + {"name": "Threshold fraction", "value": "0.5", "description": "Determines the threshold that the closed surface is created at as a fractional value between 0 and 1."} + ], + "masterRepresentation": "Binary labelmap", + "referenceImageExtentOffset": [0, 0, 0] + "segments": [ + { + "color": [0.992157, 0.909804, 0.619608], + "colorAutoGenerated": true, + "extent": [0, 124, 0, 127, 0, 33], + "id": "Segment_1", + "labelValue": 1, + "layer": 0, + "name": "ribs", + "nameAutoGenerated": true, + "status": "inprogress", + "terminology": { + "contextName": "Segmentation category and type - 3D Slicer General Anatomy list", + "category": ["SCT", "123037004", "Anatomical Structure"], + "type": ["SCT", "113197003", "Rib"] } + }, + { + "color": [1.0, 1.0, 0.811765], + "colorAutoGenerated": true, + "extent": [0, 124, 0, 127, 0, 33], + "id": "Segment_2", + "labelValue": 2, + "layer": 0, + "name": "cervical vertebral column", + "nameAutoGenerated": true, + "status": "inprogress", + "terminology": { + "contextName": "Segmentation category and type - 3D Slicer General Anatomy list", + "category": ["SCT", "123037004", "Anatomical Structure"], + "type": ["SCT", "122494005", "Cervical spine"] }, + "tags": { + "Some field": "some value" } + } + ] + } + """ + + from collections import OrderedDict + import logging + import nrrd + import numpy as np + import re + + if skip_voxels: + header = nrrd.read_header(filename) + voxels = None + else: + voxels, header = nrrd.read(filename) + + segmentation = OrderedDict() + + segments_fields = {} # map from segment index to key:value map + + multiple_layers = False + spaceToLps = np.eye(4) + ijkToSpace = np.eye(4) + + # Store header fields + for header_key in header: + if header_key in ["type", "endian", "dimension", "sizes"]: + # these are stored in the voxel array, it would be redundant to store in metadata + continue + + if header_key == "space": + if header[header_key] == "left-posterior-superior": + spaceToLps = np.eye(4) + elif header[header_key] == "right-anterior-superior": + spaceToLps = np.diag([-1.0, -1.0, 1.0, 1.0]) + else: + # LPS and RAS are the most commonly used image orientations, for now we only support these + raise IOError("space field must be 'left-posterior-superior' or 'right-anterior-superior'") + continue + elif header_key == 'kinds': + if header[header_key] == ['domain', 'domain', 'domain']: + multiple_layers = False + elif header[header_key] == ['list', 'domain', 'domain', 'domain']: + multiple_layers = True + else: + raise IOError("kinds field must be 'domain domain domain' or 'list domain domain domain'") + continue + elif header_key == "space origin": + ijkToSpace[0:3, 3] = header[header_key] + continue + elif header_key == "space directions": + space_directions = header[header_key] + if space_directions.shape[0] == 4: + # 4D segmentation, skip first (nan) row + ijkToSpace[0:3, 0:3] = header[header_key][1:4, 0:3] + else: + ijkToSpace[0:3, 0:3] = header[header_key] + continue + elif header_key == "Segmentation_ContainedRepresentationNames": + # Segmentation_ContainedRepresentationNames:=Binary labelmap|Closed surface| + representations = header[header_key].split("|") + representations[:] = [item for item in representations if item != ''] # Remove empty elements + segmentation["containedRepresentationNames"] = representations + continue + elif header_key == "Segmentation_ConversionParameters": + parameters = [] + # Segmentation_ConversionParameters:=Collapse labelmaps|1|Merge the labelmaps into as few...&Compute surface normals|1|Compute...&Crop to reference image geometry|0|Crop the model...& + parameters_str = header[header_key].split("&") + for parameter_str in parameters_str: + if not parameter_str.strip(): + # empty parameter description is ignored + continue + parameter_info = parameter_str.split("|") + if len(parameter_info) != 3: + raise IOError("Segmentation_ConversionParameters field value is invalid (each parameter must be defined by 3 strings)") + parameters.append({"name": parameter_info[0], "value": parameter_info[1], "description": parameter_info[2]}) + if parameters: + segmentation["conversionParameters"] = parameters + continue + elif header_key == "Segmentation_MasterRepresentation": + # Segmentation_MasterRepresentation:=Binary labelmap + segmentation["masterRepresentation"] = header[header_key] + continue + elif header_key == "Segmentation_ReferenceImageExtentOffset": + # Segmentation_ReferenceImageExtentOffset:=0 0 0 + segmentation["referenceImageExtentOffset"] = [int(i) for i in header[header_key].split(" ")] + continue + + segment_match = re.match("^Segment([0-9]+)_(.+)", header_key) + if segment_match: + # Store in segment_fields (segmentation field) + segment_index = int(segment_match.groups()[0]) + segment_key = segment_match.groups()[1] + if segment_index not in segments_fields: + segments_fields[segment_index] = {} + segments_fields[segment_index][segment_key] = header[header_key] + continue + + segmentation[header_key] = header[header_key] + + # Compute voxel to physical transformation matrix + ijkToLps = np.dot(spaceToLps, ijkToSpace) + segmentation["ijkToLPS"] = ijkToLps + + segmentation["voxels"] = voxels + + # Process segment_fields to build segment_info + + # Get all used segment IDs (necessary for creating unique segment IDs) + segment_ids = set() + for segment_index in segments_fields: + if "ID" in segments_fields[segment_index]: + segment_ids.add(segments_fields[segment_index]["ID"]) + + # Store segment metadata in segments_info + segments_info = [] + for segment_index in sorted(segments_fields.keys()): + segment_fields = segments_fields[segment_index] + if "ID" in segment_fields: # Segment0_ID:=Segment_1 + segment_id = segment_fields["ID"] + else: + segment_id = generate_unique_segment_id(segment_ids) + segment_ids.add(segment_id) + logging.debug(f"Segment ID was not found for index {segment_index}, use automatically generated ID: {segment_id}") + + segment_info = {} + segment_info["id"] = segment_id + if "Color" in segment_fields: + segment_info["color"] = [float(i) for i in segment_fields["Color"].split(" ")] # Segment0_Color:=0.501961 0.682353 0.501961 + if "ColorAutoGenerated" in segment_fields: + segment_info["colorAutoGenerated"] = int(segment_fields["ColorAutoGenerated"]) != 0 # Segment0_ColorAutoGenerated:=1 + if "Extent" in segment_fields: + segment_info["extent"] = [int(i) for i in segment_fields["Extent"].split(" ")] # Segment0_Extent:=68 203 53 211 24 118 + if "LabelValue" in segment_fields: + segment_info["labelValue"] = int(segment_fields["LabelValue"]) # Segment0_LabelValue:=1 + if "Layer" in segment_fields: + segment_info["layer"] = int(segment_fields["Layer"]) # Segment0_Layer:=0 + if "Name" in segment_fields: + segment_info["name"] = segment_fields["Name"] # Segment0_Name:=Segment_1 + if "NameAutoGenerated" in segment_fields: + segment_info["nameAutoGenerated"] = int(segment_fields["NameAutoGenerated"]) != 0 # Segment0_NameAutoGenerated:=1 # Segment0_Tags:=Segmentation.Status:inprogress|TerminologyEntry:Segmentation category and type - 3D Slicer General Anatomy list # ~SCT^85756007^Tissue~SCT^85756007^Tissue~^^~Anatomic codes - DICOM master list~^^~^^| - tags = {} - tags_str = header[prefix + "Tags"].split("|") - for tag_str in tags_str: - tag_str = tag_str.strip() - if not tag_str: + if "Tags" in segment_fields: + tags = {} + tags_str = segment_fields["Tags"].split("|") + for tag_str in tags_str: + tag_str = tag_str.strip() + if not tag_str: + continue + key, value = tag_str.split(":", maxsplit=1) + # Process known tags: TerminologyEntry and Segmentation.Status, store all other tags as they are + if key == "TerminologyEntry": + segment_info["terminology"] = terminology_entry_from_string(value) + elif key == "Segmentation.Status": + segment_info["status"] = value + else: + tags[key] = value + if tags: + segment_info["tags"] = tags + segments_info.append(segment_info) + + segmentation["segments"] = segments_info + + return segmentation + + +def write_segmentation(file, segmentation, compression_level=9, index_order=None): + """ + Extracts segments from a segmentation volume and header. + :param segmentation: segmentation metadata and voxels + """ + import numpy as np + + voxels = segmentation["voxels"] + if voxels is None: + raise ValueError("Segmentation does not contain voxels") + + # Copy non-segmentation fields to the extracted header + output_header = {} + + for key in segmentation: + if key == "voxels": + # written separately + continue + if key == "segments": + # written later + continue + elif key == "ijkToLPS": + # image geometry will be set later in space directions, space origin fields + ijkToLPS = segmentation[key] + continue + elif key == "containedRepresentationNames": + # Segmentation_ContainedRepresentationNames:=Binary labelmap|Closed surface| + # An extra closing "|" is added as this is requires by some older Slicer versions. + representations = "|".join(segmentation[key]) + "|" + output_header["Segmentation_ContainedRepresentationNames"] = representations + elif key == "conversionParameters": + # Segmentation_ConversionParameters:=Collapse labelmaps|1|Merge the labelmaps into as few...&Compute surface normals|1|Compute...&Crop to reference image geometry|0|Crop the model...& + parameters_str = "" + parameters = segmentation[key] + for parameter in parameters: + if parameters_str != "": + parameters_str += "&" + parameters_str += f"{parameter['name']}|{parameter['value']}|{parameter['description']}" + output_header["Segmentation_ConversionParameters"] = parameters_str + elif key == "masterRepresentation": + # Segmentation_MasterRepresentation:=Binary labelmap + output_header["Segmentation_MasterRepresentation"] = segmentation[key] + elif key == "referenceImageExtentOffset": + # Segmentation_ReferenceImageExtentOffset:=0 0 0 + offset = segmentation[key] + output_header["Segmentation_ReferenceImageExtentOffset"] = " ".join([str(i) for i in offset]) + else: + output_header[key] = segmentation[key] + + # Add kinds, space directions, space origin to the header + # kinds: list domain domain domain + kinds = ["domain", "domain", "domain"] + + # space directions: (0,1,0) (0,0,-1) (-1.2999954223632812,0,0) + # 'space directions', array([ + # [ 0. , 1. , 0. ], + # [ 0. , 0. , -1. ], + # [-1.29999542, 0. , 0. ]])) + space_directions = np.array(ijkToLPS)[0:3, 0:3] + + # Add 4th dimension metadata if array is 4-dimensional (there are overlapping segments) + dims = len(voxels.shape) + if dims == 4: + # kinds: list domain domain domain + # ('kinds', ['list', 'domain', 'domain', 'domain']) + kinds = ["list"] + kinds + # space directions: none (0,1,0) (0,0,-1) (-1.2999954223632812,0,0) + # 'space directions', array([ + # [ nan, nan, nan], + # [ 0. , 1. , 0. ], + # [ 0. , 0. , -1. ], + # [-1.29999542, 0. , 0. ]])) + nan_column = np.array() + space_directions = np.row_stack(([np.nan, np.nan, np.nan], space_directions)) + elif dims != 3: + raise ValueError("Unsupported number of dimensions: " + str(dims)) + + output_header["kinds"] = kinds + output_header["space directions"] = space_directions + output_header["space origin"] = np.array(ijkToLPS)[0:3, 3] + output_header["space"] = "left-posterior-superior" # DICOM uses LPS coordinate system + + # Set defaults + if "encoding" not in segmentation: + output_header["encoding"] = "gzip" + if "referenceImageExtentOffset" not in segmentation: + output_header["Segmentation_ReferenceImageExtentOffset"] = "0 0 0" + if "masterRepresentation" not in segmentation: + output_header["Segmentation_MasterRepresentation"] = "Binary labelmap" + + # Add segments fields to the header + + # Get list of segment IDs (needed if we need to genereate new ID) + segment_ids = set() + for output_segment_index, segment in enumerate(segmentation["segments"]): + if "id" in segment: + segment_ids.add(segment["id"]) + + for output_segment_index, segment in enumerate(segmentation["segments"]): + + # Copy all segment fields corresponding to this segment + output_tags = [] + for segment_key in segment: + if segment_key == "labelValue": + # Segment0_LabelValue:=1 + field_name = "LabelValue" + value = str(segment[segment_key]) + elif segment_key == "layer": + # Segment0_Layer:=0 + field_name = "Layer" + value = str(segment[segment_key]) + elif segment_key == "name": + # Segment0_Name:=Segment_1 + field_name = "Name" + value = segment[segment_key] + elif segment_key == "id": + # Segment0_ID:=Segment_1 + field_name = "ID" + value = segment[segment_key] + elif segment_key == "color": + # Segment0_Color:=0.501961 0.682353 0.501961 + field_name = "Color" + value = ' '.join([str(i) for i in segment[segment_key]]) + elif segment_key == "nameAutoGenerated": + # Segment0_NameAutoGenerated:=1 + field_name = "NameAutoGenerated" + value = 1 if segment[segment_key] else 0 + elif segment_key == "colorAutoGenerated": + # Segment0_ColorAutoGenerated:=1 + field_name = "ColorAutoGenerated" + value = 1 if segment[segment_key] else 0 + # Process information stored in tags, for example: + # Segment0_Tags:=Segmentation.Status:inprogress|TerminologyEntry:Segmentation category and type - 3D Slicer General Anatomy list + # ~SCT^85756007^Tissue~SCT^85756007^Tissue~^^~Anatomic codes - DICOM master list~^^~^^| + elif segment_key == "terminology": + # Terminology is stored in a tag + terminology_str = terminology_entry_to_string(segment[segment_key]) + output_tags.append(f"TerminologyEntry:{terminology_str}") + # Add tags later + continue + elif segment_key == "status": + # Segmentation status is stored in a tag + output_tags.append(f"Segmentation.Status:{segment[segment_key]}") + # Add tags later continue - key, value = tag_str.split(":", maxsplit=1) - if key == "TerminologyEntry": - segment["terminology"] = read_terminology_entry(value) - elif key == "Segmentation.Status": - segment["status"] = value + elif segment_key == "tags": + # Other tags + tags = segment[segment_key] + for tag_key in tags: + output_tags.append(f"{tag_key}:{tags[tag_key]}") + # Add tags later + continue + elif segment_key == "extent": + # Segment0_Extent:=68 203 53 211 24 118 + field_name = "Extent" + value = ' '.join([str(i) for i in segment[segment_key]]) else: - tags[key] = value - segment["tags"] = tags - segments.append(segment) - segment_index += 1 - segmentation_info["segments"] = segments - return segmentation_info + field_name = segment_key + value = segment[segment_key] + + output_header[f"Segment{output_segment_index}_{field_name}"] = value + + if "id" not in segment: + # If user has not specified ID, generate a unique one + new_segment_id = generate_unique_segment_id(segment_ids) + output_header[f"Segment{output_segment_index}_ID"] = new_segment_id + segment_ids.add(new_segment_id) + + if "layer" not in segment: + # If user has not specified layer, set it to 0 + output_header[f"Segment{output_segment_index}_Layer"] = "0" + if "extent" not in segment: + # If user has not specified extent, set it to the full extent + output_shape = voxels.shape[-3:] + output_header[f"Segment{output_segment_index}_Extent"] = f"0 {output_shape[0]-1} 0 {output_shape[1]-1} 0 {output_shape[2]-1}" -def segment_from_name(segmentation_info, segment_name): - for segment in segmentation_info["segments"]: + # Add tags + # Need to end with "|" as earlier Slicer versions require this + output_header[f"Segment{output_segment_index}_Tags"] = "|".join(output_tags) + "|" + + # Write segmentation to file + if index_order is None: + index_order = 'F' + import nrrd + nrrd.write(file, voxels, output_header, compression_level=compression_level, index_order=index_order) + + +def segment_from_name(segmentation, segment_name): + segments = segmentation["segments"] + for segment in segments: if segment_name == segment["name"]: return segment raise KeyError("segment not found by name " + segment_name) -def segment_names(segmentation_info): +def segments_from_name(segmentation, segment_name): + found_segments = [] + segments = segmentation["segments"] + for segment in segments: + if segment_name == segment["name"]: + found_segments.append(segment) + return found_segments + + +def segments_from_terminology(segmentation, terminology): + found_segments = [] + segments = segmentation["segments"] + for segment in segments: + if "terminology" in segment: + if terminology_entry_matches(segment["terminology"], terminology): + found_segments.append(segment) + return found_segments + +def terminology_code_matches(code1, code2): + # Coding scheme designator + if code1[0] != code2[0]: + return False + # Code value + if code1[1] != code2[1]: + return False + return True + +def terminology_entry_matches(terminology1, terminology2): + # Category + if not terminology_code_matches(terminology1['category'], terminology2['category']): + return False + # Type + if not terminology_code_matches(terminology1['type'], terminology2['type']): + return False + # Optional type modifier + if "typeModifier" in terminology1 and "typeModifier" in terminology2: + # Both have type modifier + if not terminology_code_matches(terminology1['typeModifier'], terminology2['typeModifier']): + return False + elif "typeModifier" in terminology1 or "typeModifier" in terminology2: + # Only one of the two has type modifier + return False + # Optional anatomic region + if "anatomicRegion" in terminology1 and "anatomicRegion" in terminology2: + # Both have anatomic region + if not terminology_code_matches(terminology1['anatomicRegion'], terminology2['anatomicRegion']): + return False + # Optional anatomic region modifier + if "anatomicRegionModifier" in terminology1 and "anatomicRegionModifier" in terminology2: + # Both have anatomic region modifier + if not terminology_code_matches(terminology1['anatomicRegionModifier'][0], terminology2['anatomicRegionModifier'][0]): + return False + elif "anatomicRegionModifier" in terminology1 or "anatomicRegionModifier" in terminology2: + # Only one of the two has anatomic region modifier + return False + elif "anatomicRegion" in terminology1 or "anatomicRegion" in terminology2: + # Only one of the two has anatomic region + return False + # Terminologies match + return True + +def segment_id_from_name(segmentation, segment_name): + segments = segmentation["segments"] + for segment in segments: + if segment_name == segments[segment_id]["name"]: + return segment_id + raise KeyError("segment_id not found by name " + segment_name) + + +def segment_names(segmentation): names = [] - for segment in segmentation_info["segments"]: + segments = segmentation["segments"] + for segment in segments: names.append(segment["name"]) return names -def extract_segments(voxels, header, segmentation_info, segment_names_to_label_values): +def extract_segments(segmentation, segment_names_to_label_values, minimalExtent=False): + """ + Extracts segments from a segmentation volume and header. + Segmentation is collapsed into a 3D volume, if there were overlapping segments then the ones listed later in the segment_names_to_label_values list will overwrite the earlier ones. + :param voxels: 3D or 4D array of voxel values + :param header: dictionary of NRRD header fields + :param segmentation_metadata: dictionary of segmentation metadata + :param segment_names_to_label_values: list of segment name to label value pairs + :param minimalExtent: if True then only the minimal extent of the segment is stored. False is recommended for compatibility with older Slicer versions. + :return: 3D array of extracted voxels, dictionary of extracted header fields + """ + + from collections import OrderedDict + import copy import numpy as np - import re + + voxels = segmentation["voxels"] + if voxels is None: + raise ValueError("Segmentation does not contain voxels") # Create empty array from last 3 dimensions (output will be flattened to a 3D array) - output_voxels = np.zeros(voxels.shape[-3:], dtype=voxels.dtype) + output_shape = voxels.shape[-3:] + output_voxels = np.zeros(output_shape, dtype=voxels.dtype) - # Copy non-segmentation fields to the extracted header - output_header = {} - for key in header.keys(): - if not re.match("^Segment[0-9]+_.+", key): - output_header[key] = header[key] + # Crete independent copy of the input image and segmentation + output_segmentation = OrderedDict() + + output_segmentation["voxels"] = output_voxels + + for key in segmentation: + if key == "voxels": + continue + elif key == "segments": + continue + else: + output_segmentation[key] = copy.deepcopy(segmentation[key]) + + output_segments = [] + output_segmentation["segments"] = output_segments # Copy extracted segments dims = len(voxels.shape) for output_segment_index, segment_name_to_label_value in enumerate(segment_names_to_label_values): - - # Copy relabeled voxel data - segment = segment_from_name(segmentation_info, segment_name_to_label_value[0]) - input_label_value = segment["labelValue"] - output_label_value = segment_name_to_label_value[1] - if dims == 3: - output_voxels[voxels == input_label_value] = output_label_value - elif dims == 4: - inputLayer = segment["layer"] - output_voxels[voxels[inputLayer, :, :, :] == input_label_value] = output_label_value + if type(segment_name_to_label_value[0]) is str: + # Find segment from terminology + segments = segments_from_name(segmentation, segment_name_to_label_value[0]) else: - raise ValueError("Voxel array dimension is invalid") + segments = segments_from_terminology(segmentation, segment_name_to_label_value[0]) + if not segments: + raise ValueError(f"Segment not found: {segment_name_to_label_value[0]}") + segment_id = segments[0]["id"] + output_segment = copy.deepcopy(segments[0]) + output_label_value = segment_name_to_label_value[1] + output_segment["labelValue"] = output_label_value + output_segment["layer"] = 0 # Output is a single layer (3D volume) - # Copy all segment fields corresponding to this segment - for key in header.keys(): - prefix = "Segment{0}_".format(segment["index"]) - matched = re.match("^" + prefix + "(.+)", key) - if not matched: - continue - field_name = matched.groups()[0] - if field_name == "LabelValue": - value = output_label_value - elif field_name == "Layer": - # Output is a single layer (3D volume) - value = 0 + unionOfAllExtents = [0, -1, 0, -1, 0, -1] + for segment in segments: + # Copy relabeled voxel data + input_label_value = segment["labelValue"] + if dims == 3: + segment_voxel_positions = np.where(output_voxels[voxels == input_label_value]) + elif dims == 4: + inputLayer = segment["layer"] + segment_voxel_positions = np.where(voxels[inputLayer, :, :, :] == input_label_value) else: - value = header[key] - output_header["Segment{0}_".format(output_segment_index) + field_name] = value + raise ValueError("Voxel array dimension is invalid") + output_voxels[segment_voxel_positions] = output_label_value + if minimalExtent: + if "extent" in segment: + extent = segment["extent"] + if _isValidExtent(extent): + # Valid extent, expand the current extent + if _isValidExtent(unionOfAllExtents): + for axis in range(3): + if extent[axis*2] < unionOfAllExtents[axis*2]: + unionOfAllExtents[axis*2] = extent[axis*2] + if extent[axis*2+1] > unionOfAllExtents[axis*2+1]: + unionOfAllExtents[axis*2+1] = extent[axis*2+1] + else: + unionOfAllExtents = extent.copy() - # Remove unnecessary 4th dimension (volume is collapsed into 3D) - if dims == 4: - # Remove "none" from "none (0,1,0) (0,0,-1) (-1.2999954223632812,0,0)" - output_header["space directions"] = output_header["space directions"][-3:, :] - # Remove "list" from "list domain domain domain" - output_header["kinds"] = output_header["kinds"][-3:] + if minimalExtent: + output_segment["extent"] = unionOfAllExtents + else: + # Use the full extent as segment extent. This is a workaround for a Slicer bug + # that used the first segment's extent when reading a layer, cropping all other segments + # that had larger extent than the first segment. + output_segment["extent"] = [0, output_shape[0]-1, 0, output_shape[1]-1, 0, output_shape[2]-1] + + output_segments.append(output_segment) + + return output_segmentation - return output_voxels, output_header +def _isValidExtent(extent): + return extent[0] <= extent[1] and extent[2] <= extent[3] and extent[4] <= extent[5] diff --git a/slicerio/tests/test_segmentation.py b/slicerio/tests/test_segmentation.py index 6917915..9ca12ce 100644 --- a/slicerio/tests/test_segmentation.py +++ b/slicerio/tests/test_segmentation.py @@ -21,12 +21,12 @@ def test_segmentation_read(self): input_segmentation_filepath = slicerio.get_testdata_file('Segmentation.seg.nrrd') - segmentation_info = slicerio.read_segmentation_info(input_segmentation_filepath) + segmentation = slicerio.read_segmentation(input_segmentation_filepath, skip_voxels=True) - number_of_segments = len(segmentation_info["segments"]) + number_of_segments = len(segmentation["segments"]) self.assertEqual(number_of_segments, 7) - segment_names = slicerio.segment_names(segmentation_info) + segment_names = slicerio.segment_names(segmentation) self.assertEqual(segment_names[0], 'ribs') self.assertEqual(segment_names[1], 'cervical vertebral column') self.assertEqual(segment_names[2], 'thoracic vertebral column') @@ -35,7 +35,7 @@ def test_segmentation_read(self): self.assertEqual(segment_names[5], 'left lung') self.assertEqual(segment_names[6], 'tissue') - segment = slicerio.segment_from_name(segmentation_info, segment_names[4]) + segment = slicerio.segment_from_name(segmentation, segment_names[4]) self.assertEqual(segment['id'], 'Segment_5') self.assertEqual(segment['name'], 'right lung') self.assertEqual(segment['nameAutoGenerated'], True) @@ -51,19 +51,127 @@ def test_segmentation_read(self): self.assertEqual(terminology['category'], ['SCT', '123037004', 'Anatomical Structure']) self.assertEqual(terminology['type'], ['SCT', '39607008', 'Lung']) self.assertEqual(terminology['typeModifier'], ['SCT', '24028007', 'Right']) - self.assertEqual('anatomicContextName' in terminology, False) + self.assertEqual(terminology['anatomicContextName'], 'Anatomic codes - DICOM master list') self.assertEqual('anatomicRegion' in terminology, False) self.assertEqual('anatomicRegionModifier' in terminology, False) - def test_segmentation_pixel_type(self): + def test_extract_segments(self): input_segmentation_filepath = slicerio.get_testdata_file('Segmentation.seg.nrrd') - voxels, header = nrrd.read(input_segmentation_filepath) - segmentation_info = slicerio.read_segmentation_info(input_segmentation_filepath) - extracted_voxels, extracted_header = slicerio.extract_segments( - voxels, header, segmentation_info, [('ribs', 1), ('right lung', 2)] + segmentation = slicerio.read_segmentation(input_segmentation_filepath) + + # Extract segments by name + extracted_segmentation_by_name = slicerio.extract_segments( + segmentation, [('ribs', 1), ('right lung', 3)] ) - self.assertEqual(extracted_voxels.dtype, voxels.dtype) + # Verify pixel type of new segmentation + self.assertEqual(extracted_segmentation_by_name["voxels"].dtype, segmentation["voxels"].dtype) + + # Extract segments by terminology + extracted_segmentation_by_terminology = slicerio.extract_segments( + segmentation, [ + # Note: intentionally using "ribs" instead of "Rib" (terminology value meaning) or "ribs" (segment name) + # to test that matching is based on terminology code value. + ({"category": ["SCT", "123037004", "Anatomical Structure"], "type": ["SCT", "113197003", "Ribs"]}, 1), + ({"category": ["SCT", "123037004", "Anatomical Structure"], "type": ["SCT", "39607008", "Lung"], "typeModifier": ["SCT", "24028007", "Right"]}, 3) + ]) + + # Compare the two segmentations + self._assert_segmentations_equal(extracted_segmentation_by_name, extracted_segmentation_by_terminology) + + def test_segmentation_write(self): + import numpy as np + import tempfile + + input_segmentation_filepath = slicerio.get_testdata_file('Segmentation.seg.nrrd') + segmentation = slicerio.read_segmentation(input_segmentation_filepath) + + # Get a temporary filename + output_segmentation_filepath = tempfile.mktemp() + '.seg.nrrd' + print("Temporary filename:", output_segmentation_filepath) + + # Write and re-read the segmentation + slicerio.write_segmentation(output_segmentation_filepath, segmentation) + segmentation_stored = slicerio.read_segmentation(output_segmentation_filepath) + + # Compare the two segmentations + self._assert_segmentations_equal(segmentation, segmentation_stored) + + # Clean up temporary file + import os + os.remove(output_segmentation_filepath) + + def test_segmentation_create(self): + import numpy as np + import tempfile + + # Create segmentation with two rectangular prisms, with label values 1 and 3 + voxels = np.zeros([100, 120, 150], dtype=np.uint8) + voxels[30:50, 20:60, 70:100] = 1 + voxels[70:90, 80:110, 60:110] = 3 + + segmentation = { + "voxels": voxels, + "ijkToLPS": [ + [ 0.5, 0., 0., 10], + [ 0., 0.5, 0., 30], + [ 0., 0., 0.8, 15], + [ 0., 0., 0., 1. ]], + "containedRepresentationNames": ["Binary labelmap", "Closed surface"], + "segments": [ + { + "labelValue": 1, + "terminology": { + "contextName": "Segmentation category and type - 3D Slicer General Anatomy list", + "category": ["SCT", "123037004", "Anatomical Structure"], + "type": ["SCT", "10200004", "Liver"] } + }, + { + "labelValue": 3, + "terminology": { + "contextName": "Segmentation category and type - 3D Slicer General Anatomy list", + "category": ["SCT", "123037004", "Anatomical Structure"], + "type": ["SCT", "39607008", "Lung"], + "typeModifier": ["SCT", "24028007", "Right"] } + }, + ] + } + + # Get a temporary filename + output_segmentation_filepath = tempfile.mktemp() + '.seg.nrrd' + print("Temporary filename:", output_segmentation_filepath) + + # Write and re-read the segmentation + slicerio.write_segmentation(output_segmentation_filepath, segmentation) + segmentation_stored = slicerio.read_segmentation(output_segmentation_filepath) + + # Compare the two segmentations + self._assert_segmentations_equal(segmentation, segmentation_stored) + + # Clean up temporary file + import os + os.remove(output_segmentation_filepath) + + def _assert_segmentations_equal(self, segmentation1, segmentation2): + """Compare segmentation1 to segmentation2. + Ignores segment attributes that are present in segmentation2 but not in segmentation1. + This is useful because when a segmentation is written to file then some extra attributes + may be added to the segments (such as extent). + """ + import numpy as np + for key in segmentation1: + if type(segmentation1[key]) == np.ndarray or type(segmentation2[key]) == np.ndarray: + self.assertTrue(np.allclose(segmentation1[key], segmentation2[key]), f"Failed for key {key}") + elif key == "segments": + for segment1, segment2 in zip(segmentation1[key], segmentation2[key]): + for segmentAttribute in segment1: + self.assertEqual(segment1[segmentAttribute], segment2[segmentAttribute], f"Failed for key {key}[{segmentAttribute}]") + else: + equal = (segmentation1[key] == segmentation2[key]) + if type(equal) == list: + self.assertTrue(all(equal), f"Failed for key {key}") + else: + self.assertTrue(equal, f"Failed for key {key}") if __name__ == '__main__': unittest.main()