Commit aa13d3b9 authored by Gabriel Couture's avatar Gabriel Couture

Write metadata function

parent 8425eb85
# coding: utf-8
# author: Gabriel Couture
from typing import Dict, Tuple
from pydicom import FileDataset
def apply_layer(
layer: Dict,
rtplan: FileDataset,
rtdose: FileDataset,
rtstruct: FileDataset) -> Tuple[FileDataset, FileDataset, FileDataset]:
return rtplan, rtdose, rtstruct
......@@ -2,18 +2,61 @@
# author: Gabriel Couture
import os
from os.path import join
from typing import Dict, Optional
import pydicom
from dicomgenerator import dao
from dicomgenerator.layer import apply_layer
def write_dicom_rt(
def write_metadata(
target_directory: str,
geometry: str,
spacing: float,
dose_gradient: str,
roi: str):
with open(join(target_directory, 'README.md'), 'w') as fh:
fh.write(
"# Metadata\n"
"This directory contains the DICOM-RT files "
"for a case presented below: "
f"- geometry: {geometry}\n"
f"- spacing: {spacing}mm\n"
f"- dose_gradient: {dose_gradient}\n"
f"- roi: {roi}\n"
)
def write_dicom_rt(
target_directory: str,
geometry: str,
spacing: float,
dose_gradient: str,
roi: str,
layer: Optional[Dict] = None):
"""Write the DICOM-RT files (RTDOSE, RTSTRUCT, RTPLAN).
Parameters
----------
target_directory
Target directory.
geometry
Geometry: {'cone', 'sphere', 'cylinder'}.
spacing
spacing: {0.2, 1, 2, 3}
dose_gradient
Dose gradient: {'AP', 'SI'}.
roi
Roi: {'axial', 'right'}.
layer
Dictionary of DICOM field to edit.
Returns
-------
Nothing
"""
# Getting Data
rtplan = dao.get_rtplan()
rtdose = dao.get_rtdose(dose_gradient, spacing)
......@@ -21,6 +64,7 @@ def write_dicom_rt(
dvh_data = dao.get_analytic_dvhs(geometry, spacing, dose_gradient, roi)
# Applying desired layers
rtplan, rtdose, rtstruct = apply_layer(layer, rtplan, rtdose, rtstruct)
# Writing results
os.makedirs(target_directory)
......@@ -32,3 +76,4 @@ def write_dicom_rt(
join(target_directory, 'DVH.csv'),
index=False
)
write_metadata(target_directory, geometry, spacing, dose_gradient, roi)
# coding: utf-8
# author: Gabriel Couture
import os
import shutil
import unittest
......@@ -14,10 +15,15 @@ A_ROI = 'axial'
class TestWriter(unittest.TestCase):
def setUp(self) -> None:
if os.path.exists(A_TARGET_DIRECTORY):
shutil.rmtree(A_TARGET_DIRECTORY)
def tearDown(self) -> None:
shutil.rmtree(A_TARGET_DIRECTORY)
# shutil.rmtree(A_TARGET_DIRECTORY)
pass
def test_givenTargetDirectoryAndASetOfValidParameters_whenWritingDicomRt_thenDirectoryWithDicomFilesAndDVHAndMetadataIsWorte(self):
def test_givenATargetDirectoryAndASetOfValidParameters_whenWritingDicomRt_thenDirectoryWithDicomFilesAndDVHAndMetadataIsWorte(self):
writer.write_dicom_rt(
A_TARGET_DIRECTORY,
A_GEOMETRY,
......@@ -26,5 +32,8 @@ class TestWriter(unittest.TestCase):
A_ROI
)
raise NotImplementedError
self.assertIn('DVH.csv', os.listdir(A_TARGET_DIRECTORY))
self.assertIn('RTPLAN-1.dcm', os.listdir(A_TARGET_DIRECTORY))
self.assertIn('RTDOSE-1.dcm', os.listdir(A_TARGET_DIRECTORY))
self.assertIn('RTSTRUCT-1.dcm', os.listdir(A_TARGET_DIRECTORY))
self.assertIn('README.md', os.listdir(A_TARGET_DIRECTORY))
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