soorena62
36 lines
942 B
Python
36 lines
942 B
Python
import nibabel as nib
|
|
from core.node import Node
|
|
from core.compatibility_result import CompatibilityResult, CompatibilityLevel
|
|
|
|
|
|
class ImageWriteRequirement:
|
|
data_type: str = "image"
|
|
|
|
def check(self, provided):
|
|
if provided.data_type != self.data_type:
|
|
return CompatibilityResult(
|
|
level=CompatibilityLevel.ERROR,
|
|
message="ImageWriter requires image input"
|
|
)
|
|
return CompatibilityResult(
|
|
level=CompatibilityLevel.OK,
|
|
message="Compatible"
|
|
)
|
|
|
|
|
|
class ImageWriter(Node):
|
|
def __init__(self, output_path: str):
|
|
self.output_path = output_path
|
|
|
|
super().__init__(
|
|
name="ImageWriter",
|
|
inputs={"image": ImageWriteRequirement()},
|
|
outputs={}
|
|
)
|
|
|
|
def run(self, image_array):
|
|
nib.save(
|
|
nib.Nifti1Image(image_array, None),
|
|
self.output_path
|
|
)
|