segmentations WIP

mlavik1 · Aug 26, 2024 · cb9c92d · cb9c92d
1 parent ac33fa6
commit cb9c92d
Show file tree

Hide file tree

Showing 4 changed files with 165 additions and 7 deletions.
diff --git a/Assets/Editor/VolumeRenderedObjectCustomInspector.cs b/Assets/Editor/VolumeRenderedObjectCustomInspector.cs
@@ -13,6 +13,7 @@ public class VolumeRenderedObjectCustomInspector : Editor, IProgressView
         private bool lightSettings = true;
         private bool otherSettings = true;
         private bool secondaryVolumeSettings = true;
+        private bool segmentationSettings = true;
         private float currentProgress = 1.0f;
         private string currentProgressDescrition = "";
         private bool progressDirty = false;
@@ -140,7 +141,7 @@ public override void OnInspectorGUI()
             }
 
             // Secondary volume
-            secondaryVolumeSettings = EditorGUILayout.Foldout(secondaryVolumeSettings, "Overlay volume");
+            secondaryVolumeSettings = EditorGUILayout.Foldout(secondaryVolumeSettings, "PET/overlay volume");
             VolumeDataset secondaryDataset = volrendObj.GetSecondaryDataset();
             TransferFunction secondaryTransferFunction = volrendObj.GetSecondaryTransferFunction();
             if (secondaryDataset == null)
@@ -167,6 +168,39 @@ public override void OnInspectorGUI()
                 }
             }
 
+            // Segmentations
+            segmentationSettings = EditorGUILayout.Foldout(segmentationSettings, "Segmentations");
+            List<SegmentationLabel> segmentationLabels = volrendObj.GetSegmentationLabels();
+            if (segmentationLabels != null && segmentationLabels.Count > 0)
+            {
+                for (int i = 0; i < segmentationLabels.Count; i++)
+                {
+                    EditorGUILayout.BeginHorizontal();
+                    SegmentationLabel segmentationlabel = segmentationLabels[i];
+                    segmentationlabel.name = EditorGUILayout.TextField(segmentationlabel.name);
+                    segmentationlabel.colour = EditorGUILayout.ColorField(segmentationlabel.colour);
+                    segmentationLabels[i] = segmentationlabel;
+                    if (GUILayout.Button("delete"))
+                    {
+                        segmentationLabels.RemoveAt(i);
+                        volrendObj.UpdateSegmentationLabels();
+                    }
+                    if (GUILayout.Button("test"))
+                    {
+                        volrendObj.UpdateSegmentationLabels();
+                    }
+                    EditorGUILayout.EndHorizontal();
+                }
+            }
+            if (GUILayout.Button("Add segmentation (NRRD, NIFTI)"))
+            {
+                ImportSegmentation(volrendObj);
+            }
+            /*if (GUILayout.Button("Add segmentation (DICOM)"))
+            {
+                ImportSegmentationDicom(volrendObj);
+            }*/
+
             // Other settings
             GUILayout.Space(10);
             otherSettings = EditorGUILayout.Foldout(otherSettings, "Other Settings");
@@ -234,5 +268,32 @@ private static async void ImportPetScanDicom(VolumeRenderedObject targetObject)
                 }
             }
         }
+
+        private static async void ImportSegmentation(VolumeRenderedObject targetObject)
+        {
+            string filePath = EditorUtility.OpenFilePanel("Select a folder to load", "", "");
+            ImageFileFormat imageFileFormat = DatasetFormatUtilities.GetImageFileFormat(filePath);
+            if (!File.Exists(filePath))
+            {
+                Debug.LogError($"File doesn't exist: {filePath}");
+                return;
+            }
+            if (imageFileFormat == ImageFileFormat.Unknown)
+            {
+                Debug.LogError($"Invalid file format: {Path.GetExtension(filePath)}");
+                return;
+            }
+
+            using (ProgressHandler progressHandler = new ProgressHandler(new EditorProgressView()))
+            {
+                progressHandler.StartStage(1.0f, "Importing segmentation dataset");
+                IImageFileImporter importer = ImporterFactory.CreateImageFileImporter(imageFileFormat);
+                Task<VolumeDataset> importTask = importer.ImportAsync(filePath);
+                await importTask;
+                progressHandler.EndStage();
+
+                targetObject.AddSegmentation(importTask.Result);
+            }
+        }
     }
 }
diff --git a/Assets/Scripts/Importing/Utilities/DatasetFormatUtilities.cs b/Assets/Scripts/Importing/Utilities/DatasetFormatUtilities.cs
@@ -12,7 +12,8 @@ public static ImageFileFormat GetImageFileFormat(string filePath)
                 case ".vasp":
                     return ImageFileFormat.VASP;
                 case ".nii":
-                    return ImageFileFormat.NIFTI;
+                case ".gz":
+                    return filePath.ToLower().EndsWith(".nii.gz") ? ImageFileFormat.NIFTI : ImageFileFormat.Unknown;
                 default:
                     return ImageFileFormat.Unknown;
             }

diff --git a/Assets/Scripts/VolumeData/VolumeDataset.cs b/Assets/Scripts/VolumeData/VolumeDataset.cs
@@ -74,6 +74,11 @@ public Texture3D GetDataTexture()
             }
         }
 
+        public void RecreateDataTexture()
+        {
+            dataTexture = AsyncHelper.RunSync<Texture3D>(() => CreateTextureInternalAsync(NullProgressHandler.instance));
+        }
+
         /// <summary>
         /// Gets the 3D data texture, containing the density values of the dataset.
         /// Will create the data texture if it does not exist, without blocking the main thread.
@@ -156,6 +161,11 @@ public float GetMaxDataValue()
             return maxDataValue;
         }
 
+        public void RecalculateBounds()
+        {
+            CalculateValueBounds(new NullProgressHandler());
+        }
+
         /// <summary>
         /// Ensures that the dataset is not too large.
         /// This is automatically called during import,

diff --git a/Assets/Scripts/VolumeObject/VolumeRenderedObject.cs b/Assets/Scripts/VolumeObject/VolumeRenderedObject.cs
@@ -1,9 +1,20 @@
-using System.Threading;
+using openDicom.Encoding;
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading;
 using System.Threading.Tasks;
 using UnityEngine;
 
 namespace UnityVolumeRendering
 {
+    [System.Serializable]
+    public struct SegmentationLabel
+    {
+        public int id;
+        public string name;
+        public Color colour;
+    }
+
     [ExecuteInEditMode]
     public class VolumeRenderedObject : MonoBehaviour
     {
@@ -40,6 +51,9 @@ public class VolumeRenderedObject : MonoBehaviour
         [SerializeField, HideInInspector]
         private TransferFunction secondaryTransferFunction;
 
+        [SerializeField, HideInInspector]
+        private List<SegmentationLabel> segmentationLabels = new List<SegmentationLabel>();
+
         // Minimum and maximum gradient threshold for lighting contribution. Values below min will be unlit, and between min and max will be partly shaded.
         [SerializeField, HideInInspector]
         private Vector2 gradientLightingThreshold = new Vector2(0.02f, 0.15f);
@@ -93,20 +107,92 @@ public VolumeDataset GetSecondaryDataset()
             return this.secondaryDataset;
         }
 
+        public void SetSecondaryDataset(VolumeDataset dataset)
+        {
+            this.secondaryDataset = dataset;
+            UpdateMaterialProperties();
+        }
+
         public TransferFunction GetSecondaryTransferFunction()
         {
             return this.secondaryTransferFunction;
         }
 
-        public void SetSecondaryDataset(VolumeDataset dataset)
+        public void SetSecondaryTransferFunction(TransferFunction tf)
         {
-            this.secondaryDataset = dataset;
+            this.secondaryTransferFunction = tf;
             UpdateMaterialProperties();
         }
 
-        public void SetSecondaryTransferFunction(TransferFunction tf)
+        public List<SegmentationLabel> GetSegmentationLabels()
         {
-            this.secondaryTransferFunction = tf;
+            return segmentationLabels;
+        }
+
+        public void AddSegmentation(VolumeDataset dataset)
+        {
+            if (dataset.data.Length != secondaryDataset.data.Length)
+            {
+                Debug.LogError("Can't add segmentation with different dimension than original dataset.");
+                return;
+            }
+
+            int segmentationId = segmentationLabels.Count > 0 ? segmentationLabels.Max(l => l.id) + 1 : 1;
+
+            if (segmentationLabels.Count == 0)
+            {
+                secondaryDataset = dataset;
+            }
+            else
+            {
+                for (int i = 0; i < secondaryDataset.data.Length; i++)
+                {
+                    secondaryDataset.data[i] = dataset.data[i] > 0.0f ? (float)segmentationId : secondaryDataset.data[i];
+                }
+                secondaryDataset.RecalculateBounds();
+                secondaryDataset.RecreateDataTexture();
+                secondaryDataset.GetDataTexture().filterMode = FilterMode.Point;
+            }
+            SegmentationLabel segmentationLabel = new SegmentationLabel();
+            segmentationLabel.id = segmentationId;
+            segmentationLabel.name = dataset.name;
+            segmentationLabel.colour = Random.ColorHSV();
+            segmentationLabels.Add(segmentationLabel);
+            UpdateSegmentationLabels();
+        }
+
+        public void UpdateSegmentationLabels()
+        {
+            if (segmentationLabels.Count == 0)
+            {
+                return;
+            }
+
+            segmentationLabels.OrderBy(l => l.id);
+            if (secondaryTransferFunction == null)
+            {
+                secondaryTransferFunction = ScriptableObject.CreateInstance<TransferFunction>();
+            }
+            secondaryTransferFunction.alphaControlPoints.Clear();
+            secondaryTransferFunction.colourControlPoints.Clear();
+            int maxSegmentationId = segmentationLabels[segmentationLabels.Count - 1].id;
+            float minDataValue = secondaryDataset.GetMinDataValue();
+            float maxDataValue = secondaryDataset.GetMaxDataValue();
+            secondaryTransferFunction.alphaControlPoints.Add(new TFAlphaControlPoint(0.0f, 0.0f));
+            secondaryTransferFunction.alphaControlPoints.Add(new TFAlphaControlPoint(1.0f, 1.0f));
+            for (int i = 0; i < segmentationLabels.Count; i++)
+            {
+                SegmentationLabel segmentationLabel = segmentationLabels[i];
+                float t = segmentationLabel.id / maxDataValue;
+                secondaryTransferFunction.colourControlPoints.Add(new TFColourControlPoint(t, segmentationLabel.colour));
+                if (i == 0)
+                {
+                    secondaryTransferFunction.alphaControlPoints.Add(new TFAlphaControlPoint(t - 0.01f, 0.0f));
+                    secondaryTransferFunction.alphaControlPoints.Add(new TFAlphaControlPoint(t, 1.0f));
+                }
+            }
+            secondaryTransferFunction.GenerateTexture();
+            secondaryTransferFunction.GetTexture().filterMode = FilterMode.Point;
             UpdateMaterialProperties();
         }