soorena62
18 lines
619 B
Python
18 lines
619 B
Python
import os
|
|
import luigi
|
|
|
|
class AllMasks(luigi.Task):
|
|
mask_dir = luigi.Parameter(default=os.path.join("data", "masks"))
|
|
|
|
def output(self):
|
|
# Merely as a signal of completion
|
|
return luigi.LocalTarget(os.path.join("artifacts", "all_masks.done"))
|
|
|
|
def run(self):
|
|
files = [os.path.join(self.mask_dir, f) for f in os.listdir(self.mask_dir) if f.endswith(".nii.gz")]
|
|
if not files:
|
|
raise FileNotFoundError("No masks found in data/masks")
|
|
# Just make the done signal.
|
|
with self.output().open("w") as f:
|
|
f.write(f"{len(files)} masks discovered")
|