Add All Folders

Dagster/Check_Matching/nodes/image_writer.py

@@ -0,0 +1,35 @@
+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
+        )