soorena62
28 lines
692 B
Python
28 lines
692 B
Python
import nibabel as nib
|
|
from core.node import Node
|
|
from core.contracts_radiomics import MaskDataContract
|
|
from pathlib import Path
|
|
|
|
|
|
class MaskReader(Node):
|
|
def __init__(self, path: str, geometry_id: str):
|
|
self.path = Path(path)
|
|
|
|
img = nib.load(str(self.path))
|
|
dim = "3D" if img.ndim == 3 else "2D"
|
|
|
|
super().__init__(
|
|
name="MaskReader",
|
|
outputs={
|
|
"mask": MaskDataContract(
|
|
modality="CT",
|
|
dim=dim,
|
|
geometry_id=geometry_id
|
|
)
|
|
}
|
|
)
|
|
|
|
def load_data(self):
|
|
img = nib.load(str(self.path))
|
|
return img.get_fdata()
|