Add All Folders

Luigi/engine/tasks_maskreg.py

@@ -0,0 +1,49 @@
+import os
+import luigi
+import SimpleITK as sitk
+from pathlib import Path
+
+class MaskRegistration(luigi.Task):
+    artifacts_dir = luigi.Parameter(default="artifacts")
+    mask_dir = luigi.Parameter(default=os.path.join("data", "masks"))
+
+    def requires(self):
+        from engine.tasks_filter import ImageFilter
+        from engine.tasks_masks import AllMasks
+        return {
+            "filter": ImageFilter(artifacts_dir=self.artifacts_dir),
+            "masks": AllMasks(mask_dir=self.mask_dir)
+        }
+
+    def output(self):
+        return luigi.LocalTarget(os.path.join(self.artifacts_dir, "mask_registered_index.txt"))
+
+    def run(self):
+        filt_index = os.path.join(self.artifacts_dir, "filtered_index.txt")
+        with open(filt_index, "r") as f:
+            filtered_paths = [line.strip() for line in f if line.strip()]
+
+        mask_files = [os.path.join(self.mask_dir, f) for f in os.listdir(self.mask_dir) if f.endswith(".nii.gz")]
+        if not filtered_paths or not mask_files:
+            raise FileNotFoundError("Missing filtered images or masks for mask_registration.")
+
+        if len(filtered_paths) == len(mask_files):
+            pairs = zip(filtered_paths, mask_files)
+        else:
+            ref_path = filtered_paths[0]
+            pairs = [(ref_path, m) for m in mask_files]
+
+        out_paths = []
+        for ref_img_path, mask_path in pairs:
+            ref_img = sitk.ReadImage(ref_img_path)
+            mask_img = sitk.ReadImage(mask_path)
+            identity = sitk.Transform(ref_img.GetDimension(), sitk.sitkIdentity)
+            resampled_mask = sitk.Resample(
+                mask_img, ref_img, identity, sitk.sitkNearestNeighbor, 0, mask_img.GetPixelID()
+            )
+            out_path = Path(self.artifacts_dir) / f"mask_registered_{Path(mask_path).name}"
+            sitk.WriteImage(resampled_mask, str(out_path))
+            out_paths.append(str(out_path))
+
+        with self.output().open("w") as f:
+            f.write("\n".join(out_paths))