From 0294eae4d36bad54262640d48a724a4ba0b219c8 Mon Sep 17 00:00:00 2001
From: Kyle Sunderland <sunderlandkyl@gmail.com>
Date: Tue, 19 Jul 2022 17:21:10 -0400
Subject: [PATCH] ENH: Add option to OrthovoltageDoseEngine to erase dose
 outside segment

---
 .../Widgets/Python/OrthovoltageDoseEngine.py  | 22 +++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/ExternalBeamPlanning/Widgets/Python/OrthovoltageDoseEngine.py b/ExternalBeamPlanning/Widgets/Python/OrthovoltageDoseEngine.py
index 1ec93fd7..06bea0e0 100644
--- a/ExternalBeamPlanning/Widgets/Python/OrthovoltageDoseEngine.py
+++ b/ExternalBeamPlanning/Widgets/Python/OrthovoltageDoseEngine.py
@@ -5,6 +5,7 @@
 from DoseEngines import AbstractScriptedDoseEngine
 from DoseEngines import OrthovoltageDoseEngineUtil
 from DoseEngines import EGSnrcUtil
+from SegmentEditorEffects import SegmentEditorMaskVolumeEffect
 
 #------------------------------------------------------------------------------
 #
@@ -223,6 +224,10 @@ def defineBeamParameters(self):
        "Z in the BEAMnrc run where the phase space source was scored (cm).\n\
       Only needed if DBS is enabled.", 20)
 
+    self.scriptedEngine.addBeamParameterLineEdit(
+      "Orthovoltage dose", "EraseOutsideSegment", "Erase outside segment:",
+      "If specified, causes the dose volume to be erased outside of the segment with the specified name.", "")
+
   #------------------------------------------------------------------------------
   #TODO: Add a path parameter type using the CTK path selector that saves the selections to Application Settings
   def savePathsInApplicationSettings(self, beamNode):
@@ -501,5 +506,22 @@ def calculateDoseUsingEngine(self, beamNode, resultDoseVolumeNode):
     logging.info( 'Result dose volume for beam ' + beamNode.GetName() + ' successfully loaded.\n'
       + '  Dose range: ({:.4f}-{:.4f})'.format(accumulate.GetMin()[0],accumulate.GetMax()[0]) )
 
+    # Erase dose outside of the specified segment
+    maskSegmentId = ""
+    maskSegmentName = self.scriptedEngine.parameter(beamNode, "EraseOutsideSegment")
+    segmentationNode = parentPlan.GetSegmentationNode()
+    if segmentationNode is None or segmentationNode.GetSegmentation() is None:
+      logging.error("Could not mask dose volume. Invalid segmentation.")
+    elif maskSegmentName != "":
+      maskSegmentId = segmentationNode.GetSegmentation().GetSegmentIdBySegmentName(maskSegmentName)
+
+    if maskSegmentId != "":
+      logging.info(f"Erasing dose outside of {maskSegmentName}[ID: {maskSegmentId}]")
+      operationMode = "FILL_OUTSIDE"
+      fillValues = [0]
+      SegmentEditorMaskVolumeEffect.maskVolumeWithSegment(segmentationNode, maskSegmentId, operationMode, fillValues, resultDoseVolumeNode, resultDoseVolumeNode)
+    else:
+      logging.error("Could not mask dose volume. Invalid segment id.")
+
     # Successful execution, no error message
     return ""