Skip to content

Commit

Permalink
Segmentations + PET overlay (#264)
Browse files Browse the repository at this point in the history 
* Segmentations and PET/CT rendering
* Fixed built-in NIFTI importer: Handle byte arrays + handle import failure
  • Loading branch information
@mlavik1
mlavik1 authored Aug 29, 2024
1 parent 59606fb commit 4c4efa1
Show file tree
Hide file tree
Showing 15 changed files with 630 additions and 102 deletions.
2 changes: 2 additions & 0 deletions Assets/3rdparty/Nifti.NET/Nifti.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down
34 changes: 27 additions & 7 deletions Assets/Editor/TransferFunctionEditorWindow.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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;
Expand All @@ -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);

Expand All @@ -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
Expand Down Expand Up @@ -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();
}

Expand Down
78 changes: 78 additions & 0 deletions Assets/Editor/Utils/EditorDatasetImportUtils.cs
View file
Original file line number Diff line number Diff line change
@@ -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());
}
}
}
}
}
148 changes: 148 additions & 0 deletions Assets/Editor/VolumeRenderedObjectCustomInspector.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@
using UnityEditor;
using System.Collections.Generic;
using System.Threading.Tasks;
using System.IO;
using UnityEngine.Events;

namespace UnityVolumeRendering
{
Expand All @@ -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;
Expand Down Expand Up @@ -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");
Expand All @@ -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]);
}
}
}
}
}
}
Loading

0 comments on commit 4c4efa1

Please sign in to comment.