Support total_mr segmentation labels

Author: mlavik1

Assets/Scripts/Segmentation/SegmentationBuilder.cs

@@ -1,10 +1,17 @@
 using System.Collections.Generic;
 using UnityEngine;
+using UnityEngine.Rendering;
 
 namespace UnityVolumeRendering
 {
  public class SegmentationBuilder
  {
+ public enum SegmentationClassType
+ {
+ TotalSegmentator,
+ TotalSegmentatorMR
+ }
+
  private static Dictionary<string, Color> segmentationColours = new Dictionary<string, Color>
  {
  { "liver", new Color(0.42f, 0.18f, 0.12f) },
@@ -30,12 +37,21 @@ public class SegmentationBuilder
  "rib_right_12", "sternum", "costal_cartilages"
  };
 
- public static List<SegmentationLabel> BuildSegmentations(VolumeDataset dataset)
+ private static List<string> totalSegmentatorMRLabels = new List<string>
+ { "spleen", "kidney_right", "kidney_left", "gallbladder", "liver", "stomach", "pancreas", "adrenal_gland_right", "adrenal_gland_left", "lung_left", "lung_right", "esophagus",
+ "small_bowel", "duodenum", "colon", "urinary_bladder", "prostate", "sacrum", "vertebrae", "intervertebral_discs", "spinal_cord", "heart", "aorta", "inferior_vena_cava",
+ "portal_vein_and_splenic_vein", "iliac_artery_left", "iliac_artery_right", "iliac_vena_left", "iliac_vena_right", "humerus_left", "humerus_right", "scapula_left",
+ "scapula_right", "clavicula_left", "clavicula_right", "femur_left", "femur_right", "hip_left", "hip_right", "gluteus_maximus_left", "gluteus_maximus_right", "gluteus_medius_left",
+ "gluteus_medius_right", "gluteus_minimus_left", "gluteus_minimus_right", "autochthon_left", "autochthon_right", "iliopsoas_left", "iliopsoas_right", "brain"
+ };
+
+ public static List<SegmentationLabel> BuildSegmentations(VolumeDataset dataset, SegmentationClassType classType = SegmentationClassType.TotalSegmentator)
  {
  List<SegmentationLabel> result = new List<SegmentationLabel>();
  int minSegmentationId = int.MaxValue;
  int maxSegmentationId = int.MinValue;
  HashSet<int> usedSegmentationIds = new HashSet<int>();
+ List<string> labelList = classType == SegmentationClassType.TotalSegmentatorMR ? totalSegmentatorMRLabels : totalSegmentatorLabels;
 
  for (int i = 0; i < dataset.data.Length; i++)
  {
@@ -58,9 +74,9 @@ public static List<SegmentationLabel> BuildSegmentations(VolumeDataset dataset)
  segmentationLabel.id = segmentationId;
  segmentationLabel.name = dataset.datasetName;
  segmentationLabel.colour = Random.ColorHSV();
- if (multiLabel && segmentationId < totalSegmentatorLabels.Count + 1)
+ if (multiLabel && segmentationId < labelList.Count + 1)
  {
- string labelName = totalSegmentatorLabels[segmentationId - 1];
+ string labelName = labelList[segmentationId - 1];
  if (segmentationColours.ContainsKey(labelName))
  {
  segmentationLabel.colour = segmentationColours[labelName];