Segmentation: Isolated rendre mode

mlavik1 · Aug 27, 2024 · 3b32b60 · 3b32b60
1 parent 66c6e2c
commit 3b32b60
Show file tree

Hide file tree

Showing 5 changed files with 132 additions and 28 deletions.
diff --git a/Assets/Editor/VolumeRenderedObjectCustomInspector.cs b/Assets/Editor/VolumeRenderedObjectCustomInspector.cs
@@ -142,9 +142,9 @@ public override void OnInspectorGUI()
 
             // Secondary volume
             secondaryVolumeSettings = EditorGUILayout.Foldout(secondaryVolumeSettings, "PET/overlay volume");
-            VolumeDataset secondaryDataset = volrendObj.GetSecondaryDataset();
+            OverlayType overlayType = volrendObj.GetOverlayType();
             TransferFunction secondaryTransferFunction = volrendObj.GetSecondaryTransferFunction();
-            if (secondaryDataset == null)
+            if (overlayType != OverlayType.Overlay)
             {
                 if (GUILayout.Button("Load PET (NRRD, NIFTI)"))
                 {
@@ -164,7 +164,7 @@ public override void OnInspectorGUI()
 
                 if (GUILayout.Button("Remove secondary volume"))
                 {
-                    volrendObj.SetSecondaryDataset(null);
+                    volrendObj.SetOverlayDataset(null);
                 }
             }
 
@@ -177,20 +177,24 @@ public override void OnInspectorGUI()
                 {
                     EditorGUILayout.BeginHorizontal();
                     SegmentationLabel segmentationlabel = segmentationLabels[i];
+                    EditorGUI.BeginChangeCheck();
                     segmentationlabel.name = EditorGUILayout.TextField(segmentationlabel.name);
                     segmentationlabel.colour = EditorGUILayout.ColorField(segmentationlabel.colour);
+                    bool changed = EditorGUI.EndChangeCheck();
                     segmentationLabels[i] = segmentationlabel;
                     if (GUILayout.Button("delete"))
                     {
-                        segmentationLabels.RemoveAt(i);
-                        volrendObj.UpdateSegmentationLabels();
+                        volrendObj.RemoveSegmentation(segmentationlabel.id);
                     }
-                    if (GUILayout.Button("test"))
+                    EditorGUILayout.EndHorizontal();
+                    if (changed)
                     {
                         volrendObj.UpdateSegmentationLabels();
                     }
-                    EditorGUILayout.EndHorizontal();
                 }
+
+                SegmentationRenderMode segmentationRendreMode = (SegmentationRenderMode)EditorGUILayout.EnumPopup("Render mode", volrendObj.GetSegmentationRenderMode());
+                volrendObj.SetSegmentationRenderMode(segmentationRendreMode);
             }
             if (GUILayout.Button("Add segmentation (NRRD, NIFTI)"))
             {
@@ -200,6 +204,10 @@ public override void OnInspectorGUI()
             {
                 ImportSegmentationDicom(volrendObj);
             }*/
+            if (GUILayout.Button("Clear segmentations"))
+            {
+                volrendObj.ClearSegmentations();
+            }
 
             // Other settings
             GUILayout.Space(10);
@@ -242,7 +250,7 @@ private static async void ImportPetScan(VolumeRenderedObject targetObject)
                 TransferFunction secondaryTransferFunction = ScriptableObject.CreateInstance<TransferFunction>();
                 secondaryTransferFunction.colourControlPoints = new List<TFColourControlPoint>() { new TFColourControlPoint(0.0f, Color.red), new TFColourControlPoint(1.0f, Color.red) };
                 secondaryTransferFunction.GenerateTexture();
-                targetObject.SetSecondaryDataset(importTask.Result);
+                targetObject.SetOverlayDataset(importTask.Result);
                 targetObject.SetSecondaryTransferFunction(secondaryTransferFunction);
             }
         }
@@ -263,7 +271,7 @@ private static async void ImportPetScanDicom(VolumeRenderedObject targetObject)
                     TransferFunction secondaryTransferFunction = ScriptableObject.CreateInstance<TransferFunction>();
                     secondaryTransferFunction.colourControlPoints = new List<TFColourControlPoint>() { new TFColourControlPoint(0.0f, Color.red), new TFColourControlPoint(1.0f, Color.red) };
                     secondaryTransferFunction.GenerateTexture();
-                    targetObject.SetSecondaryDataset(importTask.Result[0]);
+                    targetObject.SetOverlayDataset(importTask.Result[0]);
                     targetObject.SetSecondaryTransferFunction(secondaryTransferFunction);
                 }
             }

diff --git a/Assets/Scripts/Segmentation/OverlayType.cs b/Assets/Scripts/Segmentation/OverlayType.cs
@@ -0,0 +1,9 @@
+namespace UnityVolumeRendering
+{
+    public enum OverlayType
+    {
+        None,
+        Overlay,
+        Segmentation
+    }
+}
diff --git a/Assets/Scripts/Segmentation/SegmentationLabel.cs b/Assets/Scripts/Segmentation/SegmentationLabel.cs
@@ -0,0 +1,12 @@
+using UnityEngine;
+
+namespace UnityVolumeRendering
+{
+    [System.Serializable]
+    public struct SegmentationLabel
+    {
+        public int id;
+        public string name;
+        public Color colour;
+    }
+}
diff --git a/Assets/Scripts/Segmentation/SegmentationRenderMode.cs b/Assets/Scripts/Segmentation/SegmentationRenderMode.cs
@@ -0,0 +1,9 @@
+namespace UnityVolumeRendering
+{
+    [System.Serializable]
+    public enum SegmentationRenderMode
+    {
+        OverlayColour,
+        Isolate
+    }
+}
diff --git a/Assets/Scripts/VolumeObject/VolumeRenderedObject.cs b/Assets/Scripts/VolumeObject/VolumeRenderedObject.cs
@@ -7,14 +7,6 @@
 
 namespace UnityVolumeRendering
 {
-    [System.Serializable]
-    public struct SegmentationLabel
-    {
-        public int id;
-        public string name;
-        public Color colour;
-    }
-
     [ExecuteInEditMode]
     public class VolumeRenderedObject : MonoBehaviour
     {
@@ -54,6 +46,12 @@ public class VolumeRenderedObject : MonoBehaviour
         [SerializeField, HideInInspector]
         private List<SegmentationLabel> segmentationLabels = new List<SegmentationLabel>();
 
+        [SerializeField, HideInInspector]
+        private OverlayType overlayType = OverlayType.None;
+
+        [SerializeField, HideInInspector]
+        private SegmentationRenderMode segmentationRenderMode = SegmentationRenderMode.OverlayColour;
+
         // 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);
@@ -102,15 +100,9 @@ public SlicingPlane CreateSlicingPlane()
             return slicingPlaneComp;
         }
 
-        public VolumeDataset GetSecondaryDataset()
+        public OverlayType GetOverlayType()
         {
-            return this.secondaryDataset;
-        }
-
-        public void SetSecondaryDataset(VolumeDataset dataset)
-        {
-            this.secondaryDataset = dataset;
-            UpdateMaterialProperties();
+            return this.overlayType;
         }
 
         public TransferFunction GetSecondaryTransferFunction()
@@ -124,6 +116,34 @@ public void SetSecondaryTransferFunction(TransferFunction tf)
             UpdateMaterialProperties();
         }
 
+        public void SetOverlayDataset(VolumeDataset dataset)
+        {
+            if (dataset != null)
+            {
+                this.overlayType = OverlayType.Overlay;
+            }
+            else if(this.overlayType == OverlayType.Overlay)
+            {
+                this.overlayType = OverlayType.None;
+            }
+            this.secondaryDataset = dataset;
+            UpdateMaterialProperties();
+        }
+
+        public SegmentationRenderMode GetSegmentationRenderMode()
+        {
+            return segmentationRenderMode;
+        }
+
+        public void SetSegmentationRenderMode(SegmentationRenderMode mode)
+        {
+            if (mode != segmentationRenderMode)
+            {
+                segmentationRenderMode = mode;
+                UpdateMaterialProperties();
+            }
+        }
+
         public List<SegmentationLabel> GetSegmentationLabels()
         {
             return segmentationLabels;
@@ -137,6 +157,8 @@ public void AddSegmentation(VolumeDataset dataset)
                 return;
             }
 
+            overlayType = OverlayType.Segmentation;
+
             int segmentationId = segmentationLabels.Count > 0 ? segmentationLabels.Max(l => l.id) + 1 : 1;
 
             if (segmentationLabels.Count == 0)
@@ -161,10 +183,44 @@ public void AddSegmentation(VolumeDataset dataset)
             UpdateSegmentationLabels();
         }
 
+        public void RemoveSegmentation(int id)
+        {
+            int segmentationIndex = segmentationLabels.FindIndex(s => s.id == id);
+            if (segmentationIndex != -1)
+            {
+                segmentationLabels.RemoveAt(segmentationIndex);
+            }
+            else
+            {
+                Debug.LogError($"Segmentation not found: {id}");
+            }
+            for (int i = 0; i < secondaryDataset.data.Length; i++)
+            {
+                secondaryDataset.data[i] = secondaryDataset.data[i] == id ? 0 : secondaryDataset.data[i];
+            }
+            secondaryDataset.RecalculateBounds();
+            secondaryDataset.RecreateDataTexture();
+            secondaryDataset.GetDataTexture().filterMode = FilterMode.Point;
+            UpdateSegmentationLabels();
+        }
+
+        public void ClearSegmentations()
+        {
+            if (overlayType == OverlayType.Segmentation)
+            {
+                secondaryDataset = null;
+                secondaryTransferFunction = null;
+                overlayType = OverlayType.None;
+            }
+            segmentationLabels.Clear();
+            UpdateMaterialProperties();
+        }
+
         public void UpdateSegmentationLabels()
         {
             if (segmentationLabels.Count == 0)
             {
+                UpdateMaterialProperties();
                 return;
             }
 
@@ -454,16 +510,26 @@ private void UpdateMatInternal(Texture3D dataTexture, Texture3D gradientTexture,
                 meshRenderer.sharedMaterial.SetTexture("_GradientTex", gradientTexture);
             }
 
-            if (secondaryDataTexture != null)
+            if (overlayType != OverlayType.None && secondaryDataTexture != null)
             {
                 Texture2D secondaryTF = secondaryTransferFunction.GetTexture();
                 meshRenderer.sharedMaterial.SetTexture("_SecondaryDataTex", secondaryDataTexture);
                 meshRenderer.sharedMaterial.SetTexture("_SecondaryTFTex", secondaryTF);
-                meshRenderer.sharedMaterial.EnableKeyword("SECONDARY_VOLUME_ON");
+                if (overlayType == OverlayType.Segmentation && segmentationRenderMode == SegmentationRenderMode.Isolate)
+                {
+                    meshRenderer.sharedMaterial.EnableKeyword("MULTIVOLUME_ISOLATE");
+                    meshRenderer.sharedMaterial.DisableKeyword("MULTIVOLUME_OVERLAY");
+                }
+                else if(overlayType == OverlayType.Overlay)
+                {
+                    meshRenderer.sharedMaterial.EnableKeyword("MULTIVOLUME_OVERLAY");
+                    meshRenderer.sharedMaterial.DisableKeyword("MULTIVOLUME_ISOLATE");
+                }
             }
             else
             {
-                meshRenderer.sharedMaterial.DisableKeyword("SECONDARY_VOLUME_ON");
+                meshRenderer.sharedMaterial.DisableKeyword("MULTIVOLUME_OVERLAY");
+                meshRenderer.sharedMaterial.DisableKeyword("MULTIVOLUME_ISOLATE");
             }
 
             if (meshRenderer.sharedMaterial.GetTexture("_NoiseTex") == null)