mosamatic2 2.0.24__py3-none-any.whl

mosamatic2/core/tasks/segmentationnumpy2niftitask/segmentationnumpy2niftitask.py

@@ -0,0 +1,86 @@
+import os
+import numpy as np
+import SimpleITK as sitk
+from mosamatic2.core.tasks.task import Task
+from mosamatic2.core.managers.logmanager import LogManager
+from mosamatic2.core.utils import (
+    is_dicom, 
+    convert_numpy_array_to_png_image, 
+    convert_dicom_to_png_image,
+    AlbertaColorMap,
+)
+
+LOG = LogManager()
+
+
+class SegmentationNumpy2NiftiTask(Task):
+    INPUTS = ['images', 'segmentations']
+    PARAMS = []
+
+    def __init__(self, inputs, params, output, overwrite):
+        super(SegmentationNumpy2NiftiTask, self).__init__(inputs, params, output, overwrite)
+
+    def load_images(self):
+        images = []
+        for f in os.listdir(self.input('images')):
+            f_path = os.path.join(self.input('images'), f)
+            if is_dicom(f_path):
+                images.append(f_path)
+        return images
+
+    def load_segmentations(self):
+        segmentations = []
+        for f in os.listdir(self.input('segmentations')):
+            if f.endswith('.seg.npy'):
+                f_path = os.path.join(self.input('segmentations'), f)
+                segmentations.append(f_path)
+        return segmentations
+    
+    def create_pairs_of_images_and_segmentations(self, images, segmentations):
+        pairs = []
+        for segmentation in segmentations:
+            for image in images:
+                image_name = os.path.split(image)[1]
+                if image_name in os.path.split(segmentation)[1]:
+                    pairs.append((image, segmentation))
+        return pairs
+    
+    def load_segmentation_as_narray(self, segmentation):
+        narray = np.load(segmentation)
+        return narray
+    
+    def create_png_from_dicom(self, file_path):
+        convert_dicom_to_png_image(file_path, self.output())
+
+    def create_png_from_array(self, data, file_path):
+        png_file_name = os.path.split(file_path)[1] + '.png'
+        convert_numpy_array_to_png_image(
+            data, 
+            self.output(),
+            AlbertaColorMap(), 
+            png_file_name,
+            fig_width=10, fig_height=10,
+        )
+
+
+    def run(self):
+        images = self.load_images()
+        segmentations = self.load_segmentations()        
+        image_and_segmentation_pairs = self.create_pairs_of_images_and_segmentations(images, segmentations)
+        nr_steps = len(image_and_segmentation_pairs)
+        for step in range(nr_steps):
+            pair = image_and_segmentation_pairs[step]
+            # Load segmentation
+            segmentation_narray = self.load_segmentation_as_narray(pair[1])
+            segmentation_narray = segmentation_narray.astype(np.uint16)
+            segmentation_narray3d = segmentation_narray[None, ...]
+            # Load DICOM
+            image_itk = sitk.ReadImage(pair[0])
+            segmentation_itk = sitk.GetImageFromArray(segmentation_narray3d)
+            segmentation_itk.CopyInformation(image_itk)
+            segmentation_itk_name = os.path.split(pair[1])[1] + '.nii.gz'
+            segmentation_itk_path = os.path.join(self.output(), segmentation_itk_name)
+            sitk.WriteImage(segmentation_itk, segmentation_itk_path, useCompression=True)
+            self.create_png_from_array(segmentation_narray, os.path.join(self.output(), segmentation_itk_name))
+            self.create_png_from_dicom(pair[0])
+            self.set_progress(step, nr_steps)

mosamatic2/core/tasks/segmentmusclefatl3tensorflowtask/paramloader.py

@@ -0,0 +1,39 @@
+import json
+
+
+class ParamLoader:
+    def __init__(self, json_path):
+        self.update(json_path)
+
+    def save(self, json_path):
+        """"
+        Save dict to json file
+
+        Parameters
+        ----------
+        json_path : string
+            Path to save location
+        """
+        with open(json_path, 'w') as f:
+            json.dump(self.__dict__, f, indent=4)
+
+    def update(self, json_path):
+        """
+        Load parameters from json file
+
+        Parameters
+        ----------
+        json_path : string
+            Path to json file
+        """
+        with open(json_path) as f:
+            params = json.load(f)
+            self.__dict__.update(params)
+
+    @property
+    def dict(self):
+        """"
+        Give dict-like access to Params instance
+        by: 'params.dict['learning_rate']'
+        """
+        return self.__dict__

mosamatic2/core/tasks/segmentmusclefatl3tensorflowtask/segmentmusclefatl3tensorflowtask.py

@@ -0,0 +1,122 @@
+import os
+import zipfile
+import tempfile
+import numpy as np
+
+import models
+
+from mosamatic2.core.tasks.task import Task
+from mosamatic2.core.tasks.segmentmusclefatl3tensorflowtask.paramloader import ParamLoader
+from mosamatic2.core.data.multidicomimage import MultiDicomImage
+from mosamatic2.core.data.dicomimage import DicomImage
+from mosamatic2.core.managers.logmanager import LogManager
+from mosamatic2.core.utils import (
+    normalize_between,
+    get_pixels_from_dicom_object,
+    convert_labels_to_157,
+)
+
+DEVICE = 'cpu'
+L3_INDEX = 167
+LOG = LogManager()
+
+
+class SegmentMuscleFatL3TensorFlowTask(Task):
+    INPUTS = [
+        'images', 
+        'model_files'
+    ]
+    PARAMS = ['model_version']
+
+    def __init__(self, inputs, params, output, overwrite=True):
+        super(SegmentMuscleFatL3TensorFlowTask, self).__init__(inputs, params, output, overwrite)
+
+    def load_images(self):        
+        image_data = MultiDicomImage()
+        image_data.set_path(self.input('images'))
+        if image_data.load():
+            return image_data
+        raise RuntimeError('Could not load images')
+
+    def load_model_files(self):
+        model_files = []
+        for f in os.listdir(self.input('model_files')):
+            f_path = os.path.join(self.input('model_files'), f)
+            if f_path.endswith('.zip') or f_path.endswith('.json'):
+                model_files.append(f_path)
+        if len(model_files) != 3:
+            raise RuntimeError(f'Found {len(model_files)} model files. This should be 3!')
+        return model_files
+
+    def load_models_and_params(self, model_files, model_version):
+        tfLoaded = False
+        model, contour_model, params = None, None, None
+        for f_path in model_files:
+            f_name = os.path.split(f_path)[1]
+            if f_name == f'model-{str(model_version)}.zip':
+                if not tfLoaded:
+                    import tensorflow as tf
+                    tfLoaded = True
+                with tempfile.TemporaryDirectory() as model_dir_unzipped:
+                    os.makedirs(model_dir_unzipped, exist_ok=True)
+                    with zipfile.ZipFile(f_path) as zipObj:
+                        zipObj.extractall(path=model_dir_unzipped)
+                    model = tf.keras.models.load_model(model_dir_unzipped, compile=False)
+            elif f_name == f'contour_model-{str(model_version)}.zip':
+                if not tfLoaded:
+                    import tensorflow as tf
+                    tfLoaded = True
+                with tempfile.TemporaryDirectory() as contour_model_dir_unzipped:
+                    os.makedirs(contour_model_dir_unzipped, exist_ok=True)
+                    with zipfile.ZipFile(f_path) as zipObj:
+                        zipObj.extractall(path=contour_model_dir_unzipped)
+                    contour_model = tf.keras.models.load_model(contour_model_dir_unzipped, compile=False)
+            elif f_name == f'params-{model_version}.json':
+                params = ParamLoader(f_path)
+            else:
+                pass
+        return model, contour_model, params
+
+    def extract_contour(self, image, contour_model, params):
+        ct = np.copy(image)
+        ct = normalize_between(ct, params.dict['min_bound_contour'], params.dict['max_bound_contour'])
+        img2 = np.expand_dims(ct, 0)
+        img2 = np.expand_dims(img2, -1)
+        pred = contour_model.predict([img2])
+        pred_squeeze = np.squeeze(pred)
+        pred_max = pred_squeeze.argmax(axis=-1)
+        mask = np.uint8(pred_max)
+        return mask
+
+    def segment_muscle_and_fat(self, image, model):
+        img2 = np.expand_dims(image, 0)
+        img2 = np.expand_dims(img2, -1)
+        pred = model.predict([img2])
+        pred_squeeze = np.squeeze(pred)
+        pred_max = pred_squeeze.argmax(axis=-1)
+        return pred_max
+        
+    def process_file(self, image, output_dir, model, contour_model, params):
+        assert isinstance(image, DicomImage)
+        pixels = get_pixels_from_dicom_object(image.object(), normalize=True)
+        if contour_model:
+            mask = self.extract_contour(pixels, contour_model, params)
+            pixels = normalize_between(pixels, params.dict['min_bound'], params.dict['max_bound'])
+            pixels = pixels * mask
+        pixels = pixels.astype(np.float32)
+        segmentation = self.segment_muscle_and_fat(pixels, model)
+        segmentation = convert_labels_to_157(segmentation)
+        segmentation_file_name = os.path.split(image.path())[1]
+        segmentation_file_path = os.path.join(output_dir, f'{segmentation_file_name}.seg.npy')
+        np.save(segmentation_file_path, segmentation)
+
+    def run(self):
+        image_data = self.load_images()
+        model_files = self.load_model_files()
+        model_version = self.param('model_version')
+        model, contour_model, params = self.load_models_and_params(model_files, model_version)
+        images = image_data.images()
+        nr_steps = len(images)
+        for step in range(nr_steps):
+            self.process_file(images[step], self.output(), model, contour_model, params)
+            self.set_progress(step, nr_steps)

mosamatic2/core/tasks/segmentmusclefatt4pytorchtask/paramloader.py

@@ -0,0 +1,39 @@
+import json
+
+
+class ParamLoader:
+    def __init__(self, json_path):
+        self.update(json_path)
+
+    def save(self, json_path):
+        """"
+        Save dict to json file
+
+        Parameters
+        ----------
+        json_path : string
+            Path to save location
+        """
+        with open(json_path, 'w') as f:
+            json.dump(self.__dict__, f, indent=4)
+
+    def update(self, json_path):
+        """
+        Load parameters from json file
+
+        Parameters
+        ----------
+        json_path : string
+            Path to json file
+        """
+        with open(json_path) as f:
+            params = json.load(f)
+            self.__dict__.update(params)
+
+    @property
+    def dict(self):
+        """"
+        Give dict-like access to Params instance
+        by: 'params.dict['learning_rate']'
+        """
+        return self.__dict__

mosamatic2/core/tasks/segmentmusclefatt4pytorchtask/segmentmusclefatt4pytorchtask.py

@@ -0,0 +1,128 @@
+import os
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+
+import models
+
+from mosamatic2.core.tasks.task import Task
+from mosamatic2.core.utils import (
+    normalize_between,
+    get_pixels_from_dicom_object,
+    convert_labels_to_157,
+)
+from mosamatic2.core.data.multidicomimage import MultiDicomImage
+from mosamatic2.core.data.dicomimage import DicomImage
+from mosamatic2.core.tasks.segmentmusclefatt4pytorchtask.paramloader import ParamLoader
+
+DEVICE = 'cpu'
+
+
+class SegmentMuscleFatT4PyTorchTask(Task):
+    INPUTS = [
+        'images', 
+        'model_files'
+    ]
+    PARAMS = ['model_version']
+
+    def __init__(self, inputs, params, output, overwrite=True):
+        super(SegmentMuscleFatT4PyTorchTask, self).__init__(inputs, params, output, overwrite)
+
+    def load_images(self):        
+        image_data = MultiDicomImage()
+        image_data.set_path(self.input('images'))
+        if image_data.load():
+            return image_data
+        raise RuntimeError('Could not load images')
+
+    def load_model_files(self):
+        model_files = []
+        for f in os.listdir(self.input('model_files')):
+            f_path = os.path.join(self.input('model_files'), f)
+            if f_path.endswith('.pt') or f_path.endswith('.json'):
+                model_files.append(f_path)
+        if len(model_files) != 3:
+            raise RuntimeError(f'Found {len(model_files)} model files. This should be 3!')
+        return model_files
+
+    def load_models_and_params(self, model_files, model_version):
+        # First load params.json because it is needed to instantiate the models
+        params = None
+        for f_path in model_files:
+            f_name = os.path.split(f_path)[1]
+            if f_name == f'params-{str(model_version)}.json':
+                params = ParamLoader(f_path)
+                break
+        if params is None:
+            raise RuntimeError('Could not load parameters')
+        # Next, load models
+        model, contour_model = None, None
+        for f_path in model_files:
+            f_name = os.path.split(f_path)[1]
+            if f_name == f'model-{str(model_version)}.pt':
+                # model = models.AttentionUNet(params, 4).to(device=DEVICE)
+                model = models.UNet(params, 4).to(device=DEVICE)
+                model.load_state_dict(torch.load(f_path, weights_only=False, map_location=torch.device(DEVICE)))
+                model.eval()
+            elif f_name == f'contour_model-{str(model_version)}.pt':
+                contour_model = models.UNet(params, 2).to(device=DEVICE)
+                contour_model.load_state_dict(torch.load(f_path, weights_only=False, map_location=torch.device(DEVICE)))
+                contour_model.eval()
+            else:
+                pass
+        return model, contour_model, params
+    
+    def extract_contour(self, image, contour_model):
+        with torch.no_grad():
+            # Create 4D Tensor input
+            input = np.expand_dims(image, 0)
+            input = np.expand_dims(input, 0)
+            input = torch.Tensor(input)
+            input = input.to(DEVICE, dtype=torch.float)
+            # Predict
+            prediction = contour_model(input)
+            prediction = torch.argmax(prediction, axis=1)
+            prediction = prediction.squeeze()
+            prediction = prediction.detach().cpu().numpy()
+            result = image * prediction
+        return result
+    
+    def show_result(self, result):
+        fig = plt.figure()
+        plt.imshow(result, cmap='gray')
+        plt.savefig("/Users/ralph/result.png")
+    
+    def segment_muscle_and_fat(self, image, model):
+        input = np.expand_dims(image, 0)
+        input = np.expand_dims(input, 0)
+        input = torch.Tensor(input)
+        input = input.to(DEVICE, dtype=torch.float)
+        segmentation = model(input)
+        segmentation = torch.argmax(segmentation, axis=1)
+        segmentation = segmentation.squeeze()
+        segmentation = segmentation.detach().cpu().numpy()
+        return segmentation
+    
+    def process_file(self, image, output_dir, model, contour_model, params):
+        assert isinstance(image, DicomImage)
+        pixels = get_pixels_from_dicom_object(image.object(), normalize=True)
+        if contour_model:
+            pixels = normalize_between(pixels, params.dict['lower_bound'], params.dict['upper_bound'])
+            pixels = self.extract_contour(pixels, contour_model)
+        pixels = pixels.astype(np.float32)
+        segmentation = self.segment_muscle_and_fat(pixels, model)
+        segmentation = convert_labels_to_157(segmentation)
+        segmentation_file_name = os.path.split(image.path())[1]
+        segmentation_file_path = os.path.join(output_dir, f'{segmentation_file_name}.seg.npy')
+        np.save(segmentation_file_path, segmentation)
+
+    def run(self):
+        image_data = self.load_images()
+        model_files = self.load_model_files()
+        model_version = self.param('model_version')
+        model, contour_model, params = self.load_models_and_params(model_files, model_version)
+        images = image_data.images()
+        nr_steps = len(images)
+        for step in range(nr_steps):
+            self.process_file(images[step], self.output(), model, contour_model, params)
+            self.set_progress(step, nr_steps)

mosamatic2/core/tasks/selectslicefromscanstask/selectslicefromscanstask.py

@@ -0,0 +1,249 @@
+import os
+import math
+import tempfile
+import shutil
+import nibabel as nib
+import numpy as np
+import SimpleITK as sitk
+import matplotlib.pyplot as plt
+from totalsegmentator.python_api import totalsegmentator
+from mosamatic2.core.tasks.task import Task
+from mosamatic2.core.managers.logmanager import LogManager
+from mosamatic2.core.utils import load_dicom
+
+LOG = LogManager()
+
+TOTAL_SEGMENTATOR_OUTPUT_DIR = os.path.join(tempfile.gettempdir(), 'total_segmentator_output')
+TOTAL_SEGMENTATOR_TASK = 'total'
+Z_DELTA_OFFSETS = {
+    'vertebrae_L3': 0.333,
+    'vertebrae_T4': 0.5,
+}
+
+
+class SelectSliceFromScansTask(Task):
+    INPUTS = ['scans']
+    PARAMS = ['vertebra']
+
+    def __init__(self, inputs, params, output, overwrite):
+        super(SelectSliceFromScansTask, self).__init__(inputs, params, output, overwrite)
+        self._error_dir = os.path.split(self.output())[0]
+        self._error_dir = os.path.join(self._error_dir, 'selectslicefromscanstask_errors')
+        os.makedirs(self._error_dir, exist_ok=True)
+        self._error_file = os.path.join(self._error_dir, 'errors.txt')
+        with open(self._error_file, 'w') as f:
+            f.write('Errors:\n\n')
+        LOG.info(f'Error directory: {self._error_dir}')
+
+    def write_error(self, message):
+        LOG.error(message)
+        with open(self._error_file, 'a') as f:
+            f.write(message + '\n')
+
+    def load_scan_dirs(self):
+        scan_dirs = []
+        for d in os.listdir(self.input('scans')):
+            scan_dir = os.path.join(self.input('scans'), d)
+            if os.path.isdir(scan_dir):
+                scan_dirs.append(scan_dir)
+        return scan_dirs
+    
+    def read_ct_series_sitk(self, scan_dir):
+        reader = sitk.ImageSeriesReader()
+        series_ids = reader.GetGDCMSeriesIDs(scan_dir)
+        if not series_ids:
+            raise ValueError(f'No DICOM series found in {scan_dir}')
+        file_names = reader.GetGDCMSeriesFileNames(scan_dir, series_ids[0])
+        reader.SetFileNames(file_names)
+        img = reader.Execute()
+        return img
+    
+    def resample_to_reference(self, moving: sitk.Image, reference: sitk.Image, is_label: bool) -> sitk.Image:
+        interp = sitk.sitkNearestNeighbor if is_label else sitk.sitkLinear
+        return sitk.Resample(
+            moving,
+            reference,
+            sitk.Transform(),
+            interp,
+            0,  # default value
+            moving.GetPixelID()
+        )
+    
+    def centroid_x_index_from_mask(self, mask_ref: sitk.Image) -> int:
+        # mask_ref must already be on the CT grid
+        arr = sitk.GetArrayFromImage(mask_ref)  # shape: [z, y, x]
+        idx = np.argwhere(arr > 0)
+        if idx.size == 0:
+            # fallback to true mid-sagittal
+            size_x = mask_ref.GetSize()[0]
+            return size_x // 2
+        x_mean = int(round(idx[:, 2].mean()))
+        return x_mean
+    
+    def z_index_from_physical_z(self, ct_img: sitk.Image, z_phys: float, x_index: int) -> int:
+        # pick a y roughly mid (or better: y centroid of mask if you want)
+        size = ct_img.GetSize()  # (x, y, z)
+        y_index = size[1] // 2
+
+        # Convert chosen (x_index, y_index, any z_index) -> physical, then replace z
+        p = ct_img.TransformIndexToPhysicalPoint((x_index, y_index, size[2] // 2))
+        phys_point = (p[0], p[1], z_phys)
+
+        try:
+            ix, iy, iz = ct_img.TransformPhysicalPointToIndex(phys_point)
+        except RuntimeError:
+            # if z_phys is slightly out of range, clamp later
+            iz = None
+
+        if iz is None:
+            # clamp using physical bounds
+            # safest crude clamp: convert all slice centers to physical z and find closest
+            z_centers = []
+            for k in range(size[2]):
+                pk = ct_img.TransformIndexToPhysicalPoint((x_index, y_index, k))
+                z_centers.append(pk[2])
+            z_centers = np.array(z_centers)
+            iz = int(np.argmin(np.abs(z_centers - z_phys)))
+
+        iz = int(np.clip(iz, 0, size[2]-1))
+        return iz
+    
+    def plot_sagittal_with_vertebra_overlay(self, scan_dir, mask_file, z_vert_phys_mm, out_png):
+        ct = self.read_ct_series_sitk(scan_dir)
+        vert_mask = sitk.ReadImage(mask_file)
+        vert_mask_ref = self.resample_to_reference(vert_mask, ct, is_label=True)
+        x_idx = self.centroid_x_index_from_mask(vert_mask_ref)
+        z_idx = self.z_index_from_physical_z(ct, z_vert_phys_mm, x_idx)
+        ct_arr = sitk.GetArrayFromImage(ct).astype(np.float32)
+        mk_arr = sitk.GetArrayFromImage(vert_mask_ref).astype(np.uint8)
+        sag_ct = ct_arr[:, :, x_idx]
+        sag_mk = mk_arr[:, :, x_idx]
+        vmin, vmax = np.percentile(sag_ct, (1, 99))
+        sy = ct.GetSpacing()[1]  # y spacing (mm)
+        sz = ct.GetSpacing()[2]  # z spacing (mm)
+        aspect = sz / sy
+        plt.figure(figsize=(7, 9))
+        plt.imshow(sag_ct, cmap="gray", vmin=vmin, vmax=vmax, origin="lower", aspect=aspect)
+        plt.imshow(sag_mk, alpha=0.35, origin="lower", aspect=aspect)
+        plt.axhline(z_idx, linewidth=2)  # line across the vertebra axial slice position
+        plt.title("Sagittal view with vertebral mask overlay + selected vertebra axial slice")
+        plt.axis("off")
+        plt.savefig(out_png, bbox_inches="tight", dpi=200)
+
+    def extract_masks(self, scan_dir):
+        os.makedirs(TOTAL_SEGMENTATOR_OUTPUT_DIR, exist_ok=True)
+        totalsegmentator(input=scan_dir, output=TOTAL_SEGMENTATOR_OUTPUT_DIR, fast=True)
+        if not os.path.isfile(os.path.join(TOTAL_SEGMENTATOR_OUTPUT_DIR, 'vertebrae_L3.nii.gz')):
+            raise Exception(f'{scan_dir}: vertebrae_L3.nii.gz does not exist')
+        # os.system(f'TotalSegmentator -i {scan_dir} -o {TOTAL_SEGMENTATOR_OUTPUT_DIR} --fast')
+
+    def delete_total_segmentator_output(self):
+        if os.path.exists(TOTAL_SEGMENTATOR_OUTPUT_DIR):
+            shutil.rmtree(TOTAL_SEGMENTATOR_OUTPUT_DIR)
+
+    def get_z_delta_offset_for_mask(self, mask_name):
+        if mask_name not in Z_DELTA_OFFSETS.keys():
+            return None
+        return Z_DELTA_OFFSETS[mask_name]
+
+    def find_slice(self, scan_dir, vertebra):
+        if vertebra == 'L3':
+            vertebral_level = 'vertebrae_L3'
+        elif vertebra == 'T4':
+            vertebral_level = 'vertebrae_T4'
+        else:
+            self.write_error(f'{scan_dir}: Unknown vertbra {vertebra}. Options are "L3" and "T4"')
+            return None
+        # Find Z-positions DICOM images
+        z_positions = {}
+        for f in os.listdir(scan_dir):
+            f_path = os.path.join(scan_dir, f)
+            try:
+                p = load_dicom(f_path, stop_before_pixels=True)
+                if p is not None and hasattr(p, "ImagePositionPatient"):
+                    z_positions[p.ImagePositionPatient[2]] = f_path
+            except Exception as e:
+                self.write_error(f"{scan_dir}: Failed to load DICOM {f_path}: {e}")
+                break
+        if not z_positions:
+            self.write_error(f"{scan_dir}: No valid DICOM z-positions found.")
+            return None
+        # Find Z-position vertebral image
+        mask_file = os.path.join(TOTAL_SEGMENTATOR_OUTPUT_DIR, f'{vertebral_level}.nii.gz')
+        if not os.path.exists(mask_file):
+            self.write_error(f"{scan_dir}: Mask file not found: {mask_file}")
+            return None
+        try:
+            mask_obj = nib.load(mask_file)
+            mask = mask_obj.get_fdata()
+            affine_transform = mask_obj.affine
+        except Exception as e:
+            self.write_error(f"{scan_dir}: Failed to load mask {mask_file}: {e}")
+            return None
+        indexes = np.array(np.where(mask == 1))
+        if indexes.size == 0:
+            self.write_error(f"{scan_dir}: No voxels found in mask {mask_file} for {vertebral_level}")
+            return None        
+        try:
+            index_min = indexes.min(axis=1)
+            index_max = indexes.max(axis=1)
+        except ValueError as e:
+            self.write_error(f"{scan_dir}: Invalid indexes array for {vertebral_level}: {e}")
+            return None
+        world_min = nib.affines.apply_affine(affine_transform, index_min)
+        world_max = nib.affines.apply_affine(affine_transform, index_max)
+        z_direction = affine_transform[:3, 2][2]
+        if z_direction == 0:
+            self.write_error(f"{scan_dir}: Affine z-direction is zero.")
+            return None
+        z_sign = math.copysign(1, z_direction)
+        z_delta_offset = self.get_z_delta_offset_for_mask(vertebral_level)
+        if z_delta_offset is None:
+            return None
+        z_delta = 0.333 * abs(world_max[2] - world_min[2]) # This needs to be vertebra-specific perhaps
+        z_l3 = world_max[2] - z_sign * z_delta
+        # Find closest L3 image in DICOM set
+        positions = sorted(z_positions.keys())
+        closest_file = None
+        for z1, z2 in zip(positions[:-1], positions[1:]):
+            if min(z1, z2) <= z_l3 <= max(z1, z2):
+                closest_z = min(z_positions.keys(), key=lambda z: abs(z - z_l3))
+                closest_file = z_positions[closest_z]
+                LOG.info(f'Closest image: {closest_file}')
+                break
+        if closest_file is None:
+            self.write_error(f"{scan_dir}: No matching slice found.")
+        return closest_file, z_l3
+    
+    def run(self):
+        scan_dirs = self.load_scan_dirs()
+        vertebra = self.param('vertebra')
+        nr_steps = len(scan_dirs)
+        for step in range(nr_steps):
+            scan_dir = scan_dirs[step]
+            scan_name = os.path.split(scan_dir)[1]
+            errors = False
+            LOG.info(f'Processing {scan_dir}...')
+            try:
+                self.extract_masks(scan_dir)
+            except Exception as e:
+                self.write_error(f'{scan_dir}: Could not extract masks [{str(e)}]. Skipping scan...')
+                errors = True
+            if not errors:
+                file_path, z_vert = self.find_slice(scan_dir, vertebra)
+                if file_path is not None:
+                    extension = '' if file_path.endswith('.dcm') else '.dcm'
+                    target_file_path = os.path.join(self.output(), vertebra + '_' + scan_name + extension)
+                    shutil.copyfile(file_path, target_file_path)
+                    mask_file = os.path.join(TOTAL_SEGMENTATOR_OUTPUT_DIR, f'vertebrae_{vertebra}.nii.gz')
+                    output_png = os.path.join(self.output(), f"{vertebra}_{scan_name}_sagittal.png")
+                    self.plot_sagittal_with_vertebra_overlay(scan_dir, mask_file, z_vert, output_png)
+                else:
+                    self.write_error(f'{scan_dir}: Could not find slice for vertebral level: {vertebra}')
+                    errors = True
+                self.delete_total_segmentator_output()
+            if errors:
+                LOG.info(f'Copying problematic scan {scan_dir} to error directory: {self._error_dir}')
+                scan_error_dir = os.path.join(self._error_dir, scan_name)
+                shutil.copytree(scan_dir, scan_error_dir)
+            self.set_progress(step, nr_steps)