Segmentations + PET overlay

Assets/3rdparty/Nifti.NET/Nifti.cs

@@ -135,6 +135,8 @@ public float[] ToSingleArray()
                 return Array.ConvertAll<short, float>(this.Data as short[], Convert.ToSingle);
             else if(type == typeof(ushort))
                 return Array.ConvertAll<ushort, float>(this.Data as ushort[], Convert.ToSingle);
+            else if (type == typeof(byte))
+                return Array.ConvertAll<byte, float>(this.Data as byte[], Convert.ToSingle);
             else
                 return null;
         }

Assets/Editor/TransferFunctionEditorWindow.cs

@@ -11,6 +11,8 @@ public class TransferFunctionEditorWindow : EditorWindow
 
         private TransferFunctionEditor tfEditor = new TransferFunctionEditor();
 
+        private bool keepTf = false;
+
         public static void ShowWindow(VolumeRenderedObject volRendObj)
         {
             // Close all (if any) 2D TF editor windows
@@ -25,6 +27,21 @@ public static void ShowWindow(VolumeRenderedObject volRendObj)
             wnd.SetInitialPosition();
         }
 
+        public static void ShowWindow(VolumeRenderedObject volRendObj, TransferFunction transferFunction)
+        {
+            // Close all (if any) 2D TF editor windows
+            TransferFunction2DEditorWindow[] tf2dWnds = Resources.FindObjectsOfTypeAll<TransferFunction2DEditorWindow>();
+            foreach (TransferFunction2DEditorWindow tf2dWnd in tf2dWnds)
+                tf2dWnd.Close();
+
+            TransferFunctionEditorWindow wnd = (TransferFunctionEditorWindow)EditorWindow.GetWindow(typeof(TransferFunctionEditorWindow));
+            wnd.volRendObject = volRendObj;
+            wnd.tf = transferFunction;
+            wnd.keepTf = true;
+            wnd.Show();
+            wnd.SetInitialPosition();
+        }
+
         private void SetInitialPosition()
         {
             Rect rect = this.position;
@@ -48,8 +65,9 @@ private void OnGUI()
 
             if (volRendObject == null)
                 return;
-
-            tf = volRendObject.transferFunction;
+
+            if (!keepTf)
+                tf = volRendObject.transferFunction;
 
             Event currentEvent = new Event(Event.current);
 
@@ -62,7 +80,7 @@ private void OnGUI()
             Rect outerRect = new Rect(0.0f, 0.0f, contentWidth, contentHeight);
             Rect tfEditorRect = new Rect(outerRect.x + 20.0f, outerRect.y + 20.0f, outerRect.width - 40.0f, outerRect.height - 50.0f);
 
-            tfEditor.SetVolumeObject(volRendObject);
+            tfEditor.SetTarget(volRendObject.dataset, tf);
             tfEditor.DrawOnGUI(tfEditorRect);
 
             // Draw horizontal zoom slider
@@ -99,20 +117,22 @@ private void OnGUI()
                     TransferFunction newTF = TransferFunctionDatabase.LoadTransferFunction(filepath);
                     if(newTF != null)
                     {
-                        tf = newTF;
-                        volRendObject.SetTransferFunction(tf);
+                        tf.alphaControlPoints = newTF.alphaControlPoints;
+                        tf.colourControlPoints = newTF.colourControlPoints;
+                        tf.GenerateTexture();
                         tfEditor.ClearSelection();
                     }
                 }
             }
              // Clear TF
             if(GUI.Button(new Rect(tfEditorRect.x + 150.0f, tfEditorRect.y + tfEditorRect.height + 20.0f, 70.0f, 30.0f), "Clear"))
             {
-                tf = ScriptableObject.CreateInstance<TransferFunction>();
+                tf.alphaControlPoints.Clear();
+                tf.colourControlPoints.Clear();
                 tf.alphaControlPoints.Add(new TFAlphaControlPoint(0.2f, 0.0f));
                 tf.alphaControlPoints.Add(new TFAlphaControlPoint(0.8f, 1.0f));
                 tf.colourControlPoints.Add(new TFColourControlPoint(0.5f, new Color(0.469f, 0.354f, 0.223f, 1.0f)));
-                volRendObject.SetTransferFunction(tf);
+                tf.GenerateTexture();
                 tfEditor.ClearSelection();
             }
 

Assets/Editor/Utils/EditorDatasetImportUtils.cs

@@ -0,0 +1,78 @@
+using System.Collections.Generic;
+using System.IO;
+using System.Threading.Tasks;
+using System;
+using System.Linq;
+using UnityEditor;
+using UnityEngine;
+
+namespace UnityVolumeRendering
+{
+    public class EditorDatasetImportUtils
+    {
+        public static async Task<VolumeDataset[]> ImportDicomDirectoryAsync(string dir, ProgressHandler progressHandler)
+        {
+            Debug.Log("Async dataset load. Hold on.");
+
+            List<VolumeDataset> importedDatasets = new List<VolumeDataset>();
+            bool recursive = true;
+
+            // Read all files
+            IEnumerable<string> fileCandidates = Directory.EnumerateFiles(dir, "*.*", recursive ? SearchOption.AllDirectories : SearchOption.TopDirectoryOnly)
+                .Where(p => p.EndsWith(".dcm", StringComparison.InvariantCultureIgnoreCase) || p.EndsWith(".dicom", StringComparison.InvariantCultureIgnoreCase) || p.EndsWith(".dicm", StringComparison.InvariantCultureIgnoreCase));
+
+            if (!fileCandidates.Any())
+            {
+                if (UnityEditor.EditorUtility.DisplayDialog("Could not find any DICOM files",
+                    $"Failed to find any files with DICOM file extension.{Environment.NewLine}Do you want to include files without DICOM file extension?", "Yes", "No"))
+                {
+                    fileCandidates = Directory.EnumerateFiles(dir, "*.*", recursive ? SearchOption.AllDirectories : SearchOption.TopDirectoryOnly);
+                }
+            }
+
+            if (fileCandidates.Any())
+            {
+                progressHandler.StartStage(0.2f, "Loading DICOM series");
+
+                IImageSequenceImporter importer = ImporterFactory.CreateImageSequenceImporter(ImageSequenceFormat.DICOM);
+                IEnumerable<IImageSequenceSeries> seriesList = await importer.LoadSeriesAsync(fileCandidates, new ImageSequenceImportSettings { progressHandler = progressHandler });
+
+                progressHandler.EndStage();
+                progressHandler.StartStage(0.8f);
+
+                int seriesIndex = 0, numSeries = seriesList.Count();
+                foreach (IImageSequenceSeries series in seriesList)
+                {
+                    progressHandler.StartStage(1.0f / numSeries, $"Importing series {seriesIndex + 1} of {numSeries}");
+                    VolumeDataset dataset = await importer.ImportSeriesAsync(series, new ImageSequenceImportSettings { progressHandler = progressHandler });
+                    if (dataset != null)
+                    {
+                        await OptionallyDownscale(dataset);
+                        importedDatasets.Add(dataset);
+                    }
+                    seriesIndex++;
+                    progressHandler.EndStage();
+                }
+
+                progressHandler.EndStage();
+            }
+            else
+                Debug.LogError("Could not find any DICOM files to import.");
+
+            return importedDatasets.ToArray();
+        }
+
+        public static async Task OptionallyDownscale(VolumeDataset dataset)
+        {
+            if (EditorPrefs.GetBool("DownscaleDatasetPrompt"))
+            {
+                if (EditorUtility.DisplayDialog("Optional DownScaling",
+                    $"Do you want to downscale the dataset? The dataset's dimension is: {dataset.dimX} x {dataset.dimY} x {dataset.dimZ}", "Yes", "No"))
+                {
+                    Debug.Log("Async dataset downscale. Hold on.");
+                    await Task.Run(() => dataset.DownScaleData());
+                }
+            }
+        }
+    }
+}

Assets/Editor/VolumeRenderedObjectCustomInspector.cs

@@ -2,6 +2,8 @@
 using UnityEditor;
 using System.Collections.Generic;
 using System.Threading.Tasks;
+using System.IO;
+using UnityEngine.Events;
 
 namespace UnityVolumeRendering
 {
@@ -11,6 +13,8 @@ public class VolumeRenderedObjectCustomInspector : Editor, IProgressView
         private bool tfSettings = true;
         private bool lightSettings = true;
         private bool otherSettings = true;
+        private bool overlayVolumeSettings = false;
+        private bool segmentationSettings = false;
         private float currentProgress = 1.0f;
         private string currentProgressDescrition = "";
         private bool progressDirty = false;
@@ -137,6 +141,101 @@ public override void OnInspectorGUI()
                 }
             }
 
+            // Overlay volume
+            overlayVolumeSettings = EditorGUILayout.Foldout(overlayVolumeSettings, "PET/overlay volume");
+            if (overlayVolumeSettings)
+            {
+                OverlayType overlayType = volrendObj.GetOverlayType();
+                TransferFunction secondaryTransferFunction = volrendObj.GetSecondaryTransferFunction();
+                if (overlayType != OverlayType.Overlay)
+                {
+                    if (GUILayout.Button("Load PET (NRRD, NIFTI)"))
+                    {
+                        ImportImageFileDataset(volrendObj, (VolumeDataset dataset) =>
+                        {
+                            TransferFunction secondaryTransferFunction = ScriptableObject.CreateInstance<TransferFunction>();
+                            secondaryTransferFunction.colourControlPoints = new List<TFColourControlPoint>() { new TFColourControlPoint(0.0f, Color.red), new TFColourControlPoint(1.0f, Color.red) };
+                            secondaryTransferFunction.GenerateTexture();
+                            volrendObj.SetOverlayDataset(dataset);
+                            volrendObj.SetSecondaryTransferFunction(secondaryTransferFunction);
+                        });
+                    }
+                    if (GUILayout.Button("Load PET (DICOM)"))
+                    {
+                        ImportDicomDataset(volrendObj, (VolumeDataset dataset) =>
+                        {
+                            TransferFunction secondaryTransferFunction = ScriptableObject.CreateInstance<TransferFunction>();
+                            secondaryTransferFunction.colourControlPoints = new List<TFColourControlPoint>() { new TFColourControlPoint(0.0f, Color.red), new TFColourControlPoint(1.0f, Color.red) };
+                            secondaryTransferFunction.GenerateTexture();
+                            volrendObj.SetOverlayDataset(dataset);
+                            volrendObj.SetSecondaryTransferFunction(secondaryTransferFunction);
+                        });
+                    }
+                }
+                else
+                {
+                    if (GUILayout.Button("Edit overlay transfer function"))
+                    {
+                        TransferFunctionEditorWindow.ShowWindow(volrendObj, secondaryTransferFunction);
+                    }
+
+                    if (GUILayout.Button("Remove secondary volume"))
+                    {
+                        volrendObj.SetOverlayDataset(null);
+                    }
+                }
+            }
+
+            // Segmentations
+            segmentationSettings = EditorGUILayout.Foldout(segmentationSettings, "Segmentations");
+            if (segmentationSettings)
+            {
+                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];
+                        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"))
+                        {
+                            volrendObj.RemoveSegmentation(segmentationlabel.id);
+                        }
+                        EditorGUILayout.EndHorizontal();
+                        if (changed)
+                        {
+                            volrendObj.UpdateSegmentationLabels();
+                        }
+                    }
+
+                    SegmentationRenderMode segmentationRendreMode = (SegmentationRenderMode)EditorGUILayout.EnumPopup("Render mode", volrendObj.GetSegmentationRenderMode());
+                    volrendObj.SetSegmentationRenderMode(segmentationRendreMode);
+                }
+                if (GUILayout.Button("Add segmentation (NRRD, NIFTI)"))
+                {
+                    ImportImageFileDataset(volrendObj, (VolumeDataset dataset) =>
+                    {
+                        volrendObj.AddSegmentation(dataset);
+                    });
+                }
+                if (GUILayout.Button("Add segmentation (DICOM)"))
+                {
+                    ImportDicomDataset(volrendObj, (VolumeDataset dataset) =>
+                    {
+                        volrendObj.AddSegmentation(dataset);
+                    });
+                }
+                if (GUILayout.Button("Clear segmentations"))
+                {
+                    volrendObj.ClearSegmentations();
+                }
+            }
+
             // Other settings
             GUILayout.Space(10);
             otherSettings = EditorGUILayout.Foldout(otherSettings, "Other Settings");
@@ -152,5 +251,54 @@ public override void OnInspectorGUI()
                 volrendObj.SetSamplingRateMultiplier(EditorGUILayout.Slider("Sampling rate multiplier", volrendObj.GetSamplingRateMultiplier(), 0.2f, 2.0f));
             }
         }
+        private static async void ImportImageFileDataset(VolumeRenderedObject targetObject, UnityAction<VolumeDataset> onLoad)
+        {
+            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 dataset");
+                IImageFileImporter importer = ImporterFactory.CreateImageFileImporter(imageFileFormat);
+                Task<VolumeDataset> importTask = importer.ImportAsync(filePath);
+                await importTask;
+                progressHandler.EndStage();
+
+                if (importTask.Result != null)
+                {
+                    onLoad.Invoke(importTask.Result);
+                }
+            }
+        }
+
+        private static async void ImportDicomDataset(VolumeRenderedObject targetObject, UnityAction<VolumeDataset> onLoad)
+        {
+            string dir = EditorUtility.OpenFolderPanel("Select a folder to load", "", "");
+            if (Directory.Exists(dir))
+            {
+                using (ProgressHandler progressHandler = new ProgressHandler(new EditorProgressView()))
+                {
+                    progressHandler.StartStage(1.0f, "Importing dataset");
+                    Task<VolumeDataset[]> importTask = EditorDatasetImportUtils.ImportDicomDirectoryAsync(dir, progressHandler);
+                    await importTask;
+                    progressHandler.EndStage();
+
+                    if (importTask.Result.Length > 0)
+                    {
+                        onLoad.Invoke(importTask.Result[0]);
+                    }
+                }
+            }
+        }
     }
 }