From cecc2225d136f2228ee1029fd7406e633625cd11 Mon Sep 17 00:00:00 2001
From: Matthew Field <matthew.field@unsw.edu.au>
Date: Thu, 9 Nov 2023 00:32:00 +1100
Subject: [PATCH] revised get_roi_names for geometric type

---
 rt_utils/rtstruct.py | 20 ++++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/rt_utils/rtstruct.py b/rt_utils/rtstruct.py
index dfe82be..3aebd7f 100644
--- a/rt_utils/rtstruct.py
+++ b/rt_utils/rtstruct.py
@@ -88,17 +88,29 @@ def validate_mask(self, mask: np.ndarray) -> bool:
 
         return True
 
-    def get_roi_names(self) -> List[str]:
+    def get_roi_names(self, geometry: str = 'ALL') -> List[str]:
         """
         Returns a list of the names of all ROI within the RTStruct
+        Optional argument to return only ROI of specific geometric type: E.g. [CLOSED_PLANAR, POINT]
         """
 
         if not self.ds.StructureSetROISequence:
+            print("StructureSetROISequence is empty.")
             return []
 
-        return [
-            structure_roi.ROIName for structure_roi in self.ds.StructureSetROISequence
-        ]
+        if geometry == 'ALL':
+            return [
+                structure_roi.ROIName for structure_roi in self.ds.StructureSetROISequence
+            ]
+        if geometry in ['CLOSED_PLANAR', 'OPEN_PLANAR', 'OPEN_NONPLANAR', 'POINT']:
+            return [
+                    self.ds.StructureSetROISequence[i].ROIName
+                    for i in range(0,len(self.ds.ROIContourSequence))
+                    if self.ds.ROIContourSequence[i].ContourSequence[0].ContourGeometricType==geometry
+                ]
+        else:
+            print(f"Unrecognised geometric type: {geometry}")
+            return []
 
     def get_roi_mask_by_name(self, name) -> np.ndarray:
         """