Add All Folders

2026-02-08 04:38:10 +03:30
commit 9bb57e4799
191 changed files with 8469 additions and 0 deletions
@@ -0,0 +1,73 @@
+import os
+import luigi
+import SimpleITK as sitk
+from pathlib import Path
+from engine.utils import ensure_dir
+
+def cast_to_float32(img: sitk.Image) -> sitk.Image:
+    return sitk.Cast(img, sitk.sitkFloat32)
+
+def make_initial_transform(fixed: sitk.Image, moving: sitk.Image) -> sitk.Transform:
+    dim = fixed.GetDimension()
+    if dim == 2:
+        return sitk.CenteredTransformInitializer(
+            fixed, moving, sitk.Euler2DTransform(), sitk.CenteredTransformInitializerFilter.GEOMETRY
+        )
+    elif dim == 3:
+        return sitk.CenteredTransformInitializer(
+            fixed, moving, sitk.VersorRigid3DTransform(), sitk.CenteredTransformInitializerFilter.GEOMETRY
+        )
+    else:
+        raise RuntimeError(f"Unsupported image dimension: {dim}")
+
+class ImageRegistration(luigi.Task):
+    artifacts_dir = luigi.Parameter(default="artifacts")
+
+    def requires(self):
+        from engine.tasks_reader import ImageReader
+        return ImageReader(artifacts_dir=self.artifacts_dir)
+
+    def output(self):
+        return luigi.LocalTarget(os.path.join(self.artifacts_dir, "registered_index.txt"))
+
+    def run(self):
+        # Reader Path Reader
+        reader_index = os.path.join(self.artifacts_dir, "reader", "reader_index.txt")
+        with open(reader_index, "r") as f:
+            reader_paths = [line.strip() for line in f if line.strip()]
+
+        fixed_raw = sitk.ReadImage(reader_paths[0])
+        fixed = cast_to_float32(fixed_raw)
+
+        R = sitk.ImageRegistrationMethod()
+        if fixed.GetDimension() == 2:
+            R.SetMetricAsMeanSquares()
+            R.SetInterpolator(sitk.sitkLinear)
+        else:
+            R.SetMetricAsMattesMutualInformation(numberOfHistogramBins=50)
+            R.SetMetricSamplingStrategy(R.RANDOM)
+            R.SetMetricSamplingPercentage(0.2)
+            R.SetInterpolator(sitk.sitkLinear)
+
+        R.SetOptimizerAsRegularStepGradientDescent(
+            learningRate=2.0, minStep=1e-4, numberOfIterations=200, gradientMagnitudeTolerance=1e-8
+        )
+        R.SetOptimizerScalesFromPhysicalShift()
+        R.SetShrinkFactorsPerLevel(shrinkFactors=[4, 2, 1])
+        R.SetSmoothingSigmasPerLevel(smoothingSigmas=[2, 1, 0])
+        R.SmoothingSigmasAreSpecifiedInPhysicalUnitsOn()
+
+        out_paths = []
+        for img_path in reader_paths:
+            moving_raw = sitk.ReadImage(img_path)
+            moving = cast_to_float32(moving_raw)
+            init_tx = make_initial_transform(fixed, moving)
+            R.SetInitialTransform(init_tx, inPlace=False)
+            final_tx = R.Execute(fixed, moving)
+            registered = sitk.Resample(moving, fixed, final_tx, sitk.sitkLinear, 0.0, sitk.sitkFloat32)
+            out_path = Path(self.artifacts_dir) / f"registered_{Path(img_path).name}"
+            sitk.WriteImage(registered, str(out_path))
+            out_paths.append(str(out_path))
+
+        with self.output().open("w") as f:
+            f.write("\n".join(out_paths))