Commit ac0c755e authored by Gabriel Couture's avatar Gabriel Couture

Adding get_rtdose to dao

parent e581275e
......@@ -5,7 +5,7 @@ from typing import List
import pydicom
from dicomgenerator.exceptions import WrongGeometryError
from dicomgenerator.exceptions import WrongGeometryError, InvalidDoseGradientError
def _format_spacing(spacing: float) -> str:
......@@ -17,15 +17,28 @@ def _format_spacing(spacing: float) -> str:
return spacing_str
def _format_dose_gradient(dose_gradient: str) -> str:
if dose_gradient == 'AP':
return 'AntPost'
elif dose_gradient == 'SI':
return 'SupInf'
elif dose_gradient == 'all':
return ''
raise InvalidDoseGradientError("Dose Gradient must be {'AP', 'SI', 'all'}")
def get_rtstruct(geometry: str, spacing: float) -> List[pydicom.FileDataset]:
"""Get RTSTRUCT DICOM files.
Parameters
----------
geometry
Geometry : {'cones', 'spheres', 'cylinders'}.
Geometry: {'cones', 'spheres', 'cylinders'}.
spacing
Spacing.
Spacing in millimeters: {0.2, 1, 2, 3}.
Returns
-------
......@@ -38,7 +51,7 @@ def get_rtstruct(geometry: str, spacing: float) -> List[pydicom.FileDataset]:
raise WrongGeometryError("Geometry should be in {'cones', 'spheres', 'cylinders'}")
dicom_files = filter(
lambda p: f'{_format_spacing(spacing)}_' in p and 'X' not in p,
lambda i: f'{_format_spacing(spacing)}_' in i and 'X' not in i,
os.listdir(geometry_path)
)
dicom_paths = [os.path.join(geometry_path, i) for i in dicom_files]
......@@ -46,30 +59,28 @@ def get_rtstruct(geometry: str, spacing: float) -> List[pydicom.FileDataset]:
return [pydicom.dcmread(i) for i in dicom_paths]
def get_rtdose(geometry: str, spacing: float) -> List[pydicom.FileDataset]:
"""Get RTSTRUCT DICOM files.
def get_rtdose(dose_gradient: str, spacing: float) -> List[pydicom.FileDataset]:
"""Get RTDOSE DICOM files.
Parameters
----------
geometry
Geometry : {'cones', 'spheres', 'cylinders'}.
dose_gradient
Dose gradient: {'AP', 'SI', 'all'}.
AP := Anterior Posterior, SI := Superior Inferior.
spacing
Spacing : {0.2, 1.0, 2.0, 3.0}
Spacing in millimeters: {1, 2, 3}.
Returns
-------
List of DICOM file data set.
"""
directory, _ = os.path.split(__file__)
geometry_path = os.path.join(directory, 'data', 'RTSTRUCT', geometry)
if not os.path.exists(geometry_path):
raise WrongGeometryError("Geometry should be in {'cones', 'spheres', 'cylinders'}")
rtdose_path = os.path.join(directory, 'data', 'RTDOSE')
dicom_files = filter(
lambda p: f'{_format_spacing(spacing)}_' in p and 'X' not in p,
os.listdir(geometry_path)
lambda i: f'{spacing}mm'.replace('.0', '') in i and _format_dose_gradient(dose_gradient) in i,
os.listdir(rtdose_path)
)
dicom_paths = [os.path.join(geometry_path, i) for i in dicom_files]
dicom_paths = [os.path.join(rtdose_path, i) for i in dicom_files]
return [pydicom.dcmread(i) for i in dicom_paths]
......@@ -4,3 +4,7 @@
class WrongGeometryError(ValueError):
pass
class InvalidDoseGradientError(Exception):
pass
......@@ -6,15 +6,23 @@ import pydicom
from dicomgenerator import dao
A_GEOMETRY = 'cones'
A_SPACING = 1.0
A_DOSE_GRADIENTS = 'AP'
class TestDao(unittest.TestCase):
def setUp(self) -> None:
self.a_geometry = 'cones'
self.a_spacing = 0.2
def test_givenAGeometryAndASpacing_whenGettingRtstruct_thenResultIsOneRtstructFileDataSet(self):
result = dao.get_rtstruct(A_GEOMETRY, A_SPACING)
self.assertEqual(len(result), 1)
self.assertIsInstance(result[0], pydicom.FileDataset)
self.assertEqual(result[0]['Modality'].value, 'RTSTRUCT')
def test_givenAGeometryAndASpacing_whenGettingRTSTRUCT_thenResultIsOneDICOMDataSet(self):
result = dao.get_rtstruct(self.a_geometry, self.a_spacing)
def test_givenAGAndASpacing_whenGettingRtdose_thenResultIsOneRtdoseFileDataSet(self):
result = dao.get_rtdose(A_DOSE_GRADIENTS, A_SPACING)
self.assertEqual(len(result), 1)
self.assertIsInstance(result[0], pydicom.FileDataset)
self.assertEqual(result[0]['Modality'].value, 'RTDOSE')
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment