soorena62
50 lines
2.0 KiB
Python
50 lines
2.0 KiB
Python
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))
|