threads.dl_worker

  1import os
  2import shutil
  3import sys
  4import tempfile
  5from pathlib import Path
  6
  7from PyQt6.QtCore import QThread, pyqtSignal, QProcess, QObject, QCoreApplication
  8
  9from logger import get_logger
 10from utils import get_bin_path, resource_path, get_dl_python_executable, get_script_path
 11
 12log = get_logger()
 13
 14
 15class DlWorker(QObject):
 16    progressbar_update = pyqtSignal(int)  # Overall progress (0-100)
 17    file_update = pyqtSignal(str, str)  # (filename, status)
 18    log_update = pyqtSignal(str, str)  # Log messages
 19    finished = pyqtSignal(bool, str)  # (success, message)
 20    cancel_requested = pyqtSignal()
 21
 22    def __init__(self, input_files, workspace_path, has_freesurfer):
 23        super().__init__()
 24        self.input_files = input_files
 25        self.workspace_path = workspace_path
 26        self.has_freesurfer = has_freesurfer
 27
 28        self.output_dir = None
 29
 30        self.total_files = None
 31        self.processed_files = None
 32        self.failed_files = None
 33
 34        self.is_cancelled = False
 35
 36        # Progress tracking
 37        self.total_phases = 6  # Total number of phases per file
 38        self.current_file_index = 0
 39        self.current_phase = 0
 40
 41        self.current_input_file = None
 42        self.current_input_file_basename = None
 43        self.current_synthstrip_file = None
 44        self.atlas_file = get_script_path("deep_learning/atlas/T1.nii.gz")
 45        self.brats_file = get_script_path("deep_learning/atlas/BraTS-GLI-01-001.nii")
 46        self.synthstrip_model = get_script_path("deep_learning/synthstrip.1.pt")
 47
 48        # All processes for all phases
 49        self.synthstrip_process = None
 50        self.coregistration_process = None
 51        self.reorientation_process = None
 52        self.dl_preprocess = None
 53        self.dl_process = None
 54        self.dl_postprocess = None
 55
 56        try:
 57            self.nipreps_synthstrip_bin_path = get_bin_path("nipreps-synthstrip")
 58        except FileNotFoundError:
 59            log.error(str(FileNotFoundError))
 60            raise RuntimeError
 61        log.debug(f"nipreps synthstrip binary path: {self.nipreps_synthstrip_bin_path}")
 62
 63        self.python_executable = get_dl_python_executable()
 64
 65        self.cancel_requested.connect(self.cancel)
 66
 67    def update_progress(self):
 68        """Calculates and updates the progress bar based on files and phases"""
 69        if self.total_files == 0:
 70            return
 71
 72        # Calculate progress: (completed files * 6 + current file phases) / (total files * 6) * 100
 73        total_phases_completed = (self.current_file_index * self.total_phases) + self.current_phase
 74        total_phases_overall = self.total_files * self.total_phases
 75
 76        progress = int((total_phases_completed / total_phases_overall) * 100)
 77        progress = min(progress, 100)  # Ensure it doesn't exceed 100
 78
 79        self.progressbar_update.emit(progress)
 80
 81        # Debug log (optional)
 82        self.log_update.emit(
 83            QCoreApplication.translate("DlWorker", "Progress: {0}% (File {1}/{2}, "
 84            "Phase {3}/{4})").format(progress, self.current_file_index + 1, self.total_files, self.current_phase, self.total_phases),
 85            'd'
 86        )
 87
 88    def start(self):
 89        """Processes all NIfTI files with the Deep Learning pipeline"""
 90
 91        self.total_files = len(self.input_files)
 92        log.debug(self.input_files)
 93        self.processed_files = 0
 94        self.failed_files = []
 95
 96        # Initialize progress bar
 97        self.current_file_index = 0
 98        self.current_phase = 0
 99        self.update_progress()
100        self.progressbar_update.emit(5)
101
102        self.process_single_file()
103
104    def process_single_file(self):
105        """Process single file with the Deep Learning pipeline"""
106        temp_dir = tempfile.mkdtemp(prefix=f"dl_processing_{self.current_file_index+1}_")
107        # out_dir = os.path.join(self.workspace_path, "outputs")
108        # os.makedirs(out_dir, exist_ok=True)
109        self.output_dir = temp_dir
110
111        self.current_input_file = self.input_files[self.current_file_index]
112        self.current_input_file_basename = os.path.basename(self.current_input_file)
113
114        self.log_update.emit(QCoreApplication.translate("DlWorker", "=== PROCESSING: {0} ===").format(self.current_input_file_basename), 'i')
115
116        # PHASE 1: SynthStrip
117        self.current_phase = 1
118        self.run_synthstrip()
119
120    def run_synthstrip(self):
121        """Synthstrip on a single file"""
122        phase = "Synthstrip"
123        self.file_update.emit(self.current_input_file_basename, QCoreApplication.translate("DlWorker", "Phase 1/6: Synthstrip skull strip..."))
124        self.log_update.emit(QCoreApplication.translate("DlWorker", "PHASE 1: Skull strip with Synthstrip"), 'i')
125
126        self.log_update.emit(QCoreApplication.translate("DlWorker", "SynthStrip started: {0}").format(os.path.basename(self.current_input_file)), 'i')
127
128        # Skull stripped file name
129        base_name = self.current_input_file_basename.replace(".nii.gz", "").replace(".nii", "")
130        self.current_synthstrip_file = os.path.join(self.output_dir, f"{base_name}_skull_stripped.nii.gz")
131
132        self.synthstrip_process = QProcess()
133        self.synthstrip_process.finished.connect(self.on_synthstrip_finished)
134        self.synthstrip_process.errorOccurred.connect(lambda error, string=phase: self.on_error(string, error))
135        self.synthstrip_process.readyReadStandardOutput.connect(lambda string=phase: self.on_stdout(string, self.synthstrip_process.readAllStandardOutput()))
136        self.synthstrip_process.readyReadStandardError.connect(lambda string=phase: self.on_stderr(string, self.synthstrip_process.readAllStandardError()))
137
138        if self.has_freesurfer:
139            cmd = [
140                "mri_synthstrip",
141                "-i", self.current_input_file,
142                "-o", self.current_synthstrip_file
143            ]
144        else:
145            cmd = [
146                self.nipreps_synthstrip_bin_path,
147                "-i", self.current_input_file,
148                "-o", self.current_synthstrip_file,
149                "-g",
150                "--model", self.synthstrip_model
151            ]
152
153        self.synthstrip_process.start(cmd[0], cmd[1:])
154
155    def on_synthstrip_finished(self, exit_code, exit_status):
156        if self.is_cancelled:
157            return
158
159        # This slot is called by the QProcess when it finishes
160        if exit_code != 0 or exit_status != QProcess.ExitStatus.NormalExit:
161            self.log_update.emit("SynthStrip failed", f"Exit code: {exit_code}")
162            if self.current_file_index + 1 < self.total_files:
163                self.current_file_index += 1
164                self.file_update.emit(self.current_input_file_basename, QCoreApplication.translate("DlWorker", "Segmentation failed for this file"))
165                self.process_single_file()
166            return
167
168        self.log_update.emit(QCoreApplication.translate("DlWorker", "✓ Skull stripping completed"), 'i')
169        self.update_progress()
170        # PHASE 2: Coregistration
171        self.current_phase = 2
172        self.run_coregistration()  # Start the next phase
173
174    def run_coregistration(self):
175        """Performs coregistration with atlas"""
176        phase = "Coregistration"
177
178        self.file_update.emit(self.current_input_file_basename, QCoreApplication.translate("DlWorker", "Phase 2/6: Coregistration..."))
179        self.log_update.emit(QCoreApplication.translate("DlWorker", "PHASE 2: Coregistration"), 'i')
180
181        self.log_update.emit(QCoreApplication.translate("DlWorker", "Coregistration started: {0}").format(os.path.basename(self.current_input_file)), 'i')
182
183        # Create directory for coregistration
184        coreg_dir = os.path.join(self.output_dir, "coregistration")
185        os.makedirs(coreg_dir, exist_ok=True)
186
187        self.coregistration_process = QProcess()
188        self.coregistration_process.finished.connect(self.on_coregistration_finished)
189        self.coregistration_process.errorOccurred.connect(lambda error, string=phase: self.on_error(string, error))
190        self.coregistration_process.readyReadStandardOutput.connect(lambda string=phase: self.on_stdout(string, self.coregistration_process.readAllStandardOutput()))
191        self.coregistration_process.readyReadStandardError.connect(lambda string=phase: self.on_stderr(string, self.coregistration_process.readAllStandardError()))
192
193        args = [
194            get_script_path("deep_learning/coregistration.py"),
195            "--mri", self.current_input_file,
196            "--skull", self.current_synthstrip_file,
197            "--atlas", self.atlas_file,
198            "-o", coreg_dir
199        ]
200
201        self.coregistration_process.start(self.python_executable, args)
202
203    def on_coregistration_finished(self, exit_code, exit_status):
204        if self.is_cancelled:
205            return
206
207        # This slot is called by the QProcess when it finishes
208        if exit_code != 0 or exit_status != QProcess.ExitStatus.NormalExit:
209            self.log_update.emit("Coregistration failed", f"Exit code: {exit_code}")
210            if self.current_file_index + 1 < self.total_files:
211                self.current_file_index += 1
212                self.file_update.emit(self.current_input_file_basename, QCoreApplication.translate("DlWorker", "Segmentation failed for this file"))
213                self.process_single_file()
214            return
215
216        self.log_update.emit(QCoreApplication.translate("DlWorker", "✓ Coregistration completed"), 'i')
217        self.update_progress()
218        # PHASE 3: Reorientation
219        self.current_phase = 3
220        self.run_reorientation()  # Start the next phase
221
222    def run_reorientation(self):
223        """Performs reorientation of the brain_in_atlas file using the BraTS affine matrix"""
224        phase = "Reorientation"
225
226        self.file_update.emit(self.current_input_file_basename, QCoreApplication.translate("DlWorker", "Phase 3/6: Reorientation..."))
227        self.log_update.emit(QCoreApplication.translate("DlWorker", "PHASE 3: Reorientation"), 'i')
228
229        self.log_update.emit(QCoreApplication.translate("DlWorker", "Reorientation started: {0}").format(os.path.basename(self.current_input_file)), 'i')
230
231        coreg_dir = Path(os.path.join(self.output_dir, "coregistration"))
232        brain_in_atlas_files = list(coreg_dir.glob("*.nii.gz")) + list(coreg_dir.glob("*.nii"))
233
234        if not brain_in_atlas_files:
235            self.log_update.emit(QCoreApplication.translate("DlWorker", "✗ No brain_in_atlas file found"), 'e')
236            return False
237
238        rsl_files = [f for f in brain_in_atlas_files if f.name.endswith("_rsl.nii") or f.name.endswith("_rsl.nii.gz")]
239
240        if rsl_files:
241            brain_in_atlas_file = str(rsl_files[0])
242        else:
243            # fallback
244            brain_in_atlas_file = str(brain_in_atlas_files[0])
245
246        self.reorientation_process = QProcess()
247        self.reorientation_process.finished.connect(self.on_reorientation_finished)
248        self.reorientation_process.errorOccurred.connect(lambda error, string=phase: self.on_error(string, error))
249        self.reorientation_process.readyReadStandardOutput.connect(lambda string=phase: self.on_stdout(string, self.reorientation_process.readAllStandardOutput()))
250        self.reorientation_process.readyReadStandardError.connect(lambda string=phase: self.on_stderr(string, self.reorientation_process.readAllStandardError()))
251
252        args = [
253            get_script_path("deep_learning/reorientation.py"),
254            "--input", brain_in_atlas_file,
255            "--output", self.output_dir + "/reoriented",
256            "--brats", self.brats_file,
257            "--basename", self.current_input_file_basename.replace(".nii.gz", "").replace(".nii", "")
258        ]
259
260        self.reorientation_process.start(self.python_executable, args)
261
262    def on_reorientation_finished(self, exit_code, exit_status):
263        if self.is_cancelled:
264            return
265
266        # This slot is called by the QProcess when it finishes
267        if exit_code != 0 or exit_status != QProcess.ExitStatus.NormalExit:
268            self.log_update.emit("Reorientation failed", f"Exit code: {exit_code}")
269            if self.current_file_index + 1 < self.total_files:
270                self.current_file_index += 1
271                self.file_update.emit(self.current_input_file_basename, QCoreApplication.translate("DlWorker", "Segmentation failed for this file"))
272                self.process_single_file()
273            return
274
275        self.log_update.emit(QCoreApplication.translate("DlWorker", "✓ Reorientation completed"), 'i')
276        self.update_progress()
277        # PHASE 4: Preprocess
278        self.current_phase = 4
279        self.run_preprocess()  # Start the next phase
280
281    def run_preprocess(self):
282        """Runs PHASE 4: PREPARE and PREPROCESS"""
283        phase = "Preprocessing"
284
285        self.file_update.emit(self.current_input_file_basename, QCoreApplication.translate("DlWorker", "Phase 4/6: Preparing and preprocessing..."))
286        self.log_update.emit(QCoreApplication.translate("DlWorker", "PHASE 4: Prepare and preprocess"), 'i')
287
288        data_path = os.path.join(self.output_dir, "reoriented")
289        results_path = os.path.join(self.output_dir, "preprocess")
290
291        self.dl_preprocess = QProcess()
292        self.dl_preprocess.finished.connect(self.on_preprocess_finished)
293        self.dl_preprocess.errorOccurred.connect(lambda error, string=phase: self.on_error(string, error))
294        self.dl_preprocess.readyReadStandardOutput.connect(lambda string=phase: self.on_stdout(string, self.dl_preprocess.readAllStandardOutput()))
295        self.dl_preprocess.readyReadStandardError.connect(lambda string=phase: self.on_stderr(string, self.dl_preprocess.readAllStandardError()))
296
297        args = [
298            get_script_path("deep_learning/preprocess.py"),
299            '--data', data_path,
300            '--results', results_path,
301            '--ohe'
302        ]
303
304        self.dl_preprocess.start(self.python_executable, args)
305
306    def on_preprocess_finished(self, exit_code, exit_status):
307        # This slot is called by the QProcess when it finishes
308        if self.is_cancelled:
309            return
310
311        if exit_code != 0 or exit_status != QProcess.ExitStatus.NormalExit:
312            self.log_update.emit("Preprocess failed", f"Exit code: {exit_code}")
313            if self.current_file_index + 1 < self.total_files:
314                self.current_file_index += 1
315                self.file_update.emit(self.current_input_file_basename, QCoreApplication.translate("DlWorker", "Segmentation failed for this file"))
316                self.process_single_file()
317            return
318
319        self.log_update.emit(QCoreApplication.translate("DlWorker", "✓ Preprocess completed"), 'i')
320        self.update_progress()
321        # PHASE 5: Deep learning execution
322        self.current_phase = 5
323        self.run_deep_learning()  # Start the next phase
324
325    def run_deep_learning(self):
326        """Runs PHASE 5: DEEP LEARNING"""
327        phase = "Deep Learning execution"
328
329        self.file_update.emit(self.current_input_file_basename, QCoreApplication.translate("DlWorker", "Phase 5/6: Deep Learning..."))
330        self.log_update.emit(QCoreApplication.translate("DlWorker", "PHASE 5: Deep learning execution"), 'i')
331
332        self.dl_process = QProcess()
333        self.dl_process.finished.connect(self.on_dl_finished)
334        self.dl_process.errorOccurred.connect(lambda error, string=phase: self.on_error(string, error))
335        self.dl_process.readyReadStandardOutput.connect(lambda string=phase: self.on_stdout(string, self.dl_process.readAllStandardOutput()))
336        self.dl_process.readyReadStandardError.connect(lambda string=phase: self.on_stderr(string, self.dl_process.readAllStandardError()))
337
338        args = [
339            get_script_path("deep_learning/deep_learning_runner.py"),
340            '--depth', '6',
341            '--filters', '64', '96', '128', '192', '256', '384', '512',
342            '--min_fmap', '2',
343            '--gpus', '1',
344            '--amp',
345            '--save_preds',
346            '--exec_mode', 'predict',
347            '--data', f'{self.output_dir}/preprocess/val_3d/test',
348            '--ckpt_path', get_script_path('deep_learning/checkpoints/fold3/epoch=146-dice=88.05.ckpt'),
349            '--tta',
350            '--results', f'{self.output_dir}/dl_results'
351        ]
352
353        self.dl_process.start(self.python_executable, args)
354
355    def on_dl_finished(self, exit_code, exit_status):
356        if self.is_cancelled:
357            return
358
359        # This slot is called by the QProcess when it finishes
360        if exit_code != 0 or exit_status != QProcess.ExitStatus.NormalExit:
361            self.log_update.emit("Deep learning execution failed", f"Exit code: {exit_code}")
362            if self.current_file_index + 1 < self.total_files:
363                self.current_file_index += 1
364                self.file_update.emit(self.current_input_file_basename, QCoreApplication.translate("DlWorker", "Segmentation failed for this file"))
365                self.process_single_file()
366            return
367
368        self.log_update.emit(QCoreApplication.translate("DlWorker", "✓ Deep learning execution completed"), 'i')
369        self.update_progress()
370        # PHASE 6: Postprocess
371        self.current_phase = 6
372        self.run_postprocess()  # Start the next phase
373
374    def run_postprocess(self):
375        phase = "Postprocessing"
376
377        self.file_update.emit(self.current_input_file_basename, QCoreApplication.translate("DlWorker", "Phase 6/6: Postprocessing..."))
378        self.log_update.emit(QCoreApplication.translate("DlWorker", "PHASE 6: Postprocess"), 'i')
379
380        self.dl_postprocess = QProcess()
381        self.dl_postprocess.finished.connect(self.on_postprocess_finished)
382        self.dl_postprocess.errorOccurred.connect(lambda error, string=phase: self.on_error(string, error))
383        self.dl_postprocess.readyReadStandardOutput.connect(lambda string=phase: self.on_stdout(string, self.dl_postprocess.readAllStandardOutput()))
384        self.dl_postprocess.readyReadStandardError.connect(lambda string=phase: self.on_stderr(string, self.dl_postprocess.readAllStandardError()))
385
386        args = [
387            get_script_path("deep_learning/postprocess.py"),
388            '-i', f'{self.output_dir}/dl_results/predictions_epoch=146-dice=88_05_task=train_fold=0_tta',
389            '-o', f'{self.output_dir}/dl_postprocess',
390            '--w', f'{self.workspace_path}',
391            '--atlas', self.atlas_file,
392            "--brats", self.brats_file,
393            '--mri', f'{self.current_input_file}'
394        ]
395
396        self.dl_postprocess.start(self.python_executable, args)
397
398    def on_postprocess_finished(self, exit_code, exit_status):
399        if self.is_cancelled:
400            return
401
402        # This slot is called by the QProcess when it finishes
403        if exit_code != 0 or exit_status != QProcess.ExitStatus.NormalExit:
404            self.log_update.emit("Postprocess failed", f"Exit code: {exit_code}")
405            if self.current_file_index + 1 < self.total_files:
406                self.current_file_index += 1
407                self.file_update.emit(self.current_input_file_basename, QCoreApplication.translate("DlWorker", "Segmentation failed for this file"))
408                self.process_single_file()
409            return
410
411        self.log_update.emit(QCoreApplication.translate("DlWorker", "✓ Postprocess completed"), 'i')
412        self.file_update.emit(self.current_input_file_basename,
413                              QCoreApplication.translate("DlWorker", "Segmentation completed"))
414        self.update_progress()
415
416        if self.current_file_index + 1 < self.total_files:
417            self.current_file_index += 1
418            self.process_single_file()
419        else:
420            self.finished.emit(True, QCoreApplication.translate("DlWorker", "Processing completed"))
421        return
422
423    def cancel(self):
424        self.is_cancelled = True
425        self.log_update.emit(QCoreApplication.translate("DlWorker", "Cancellation requested - stopping all processes..."), 'w')
426
427        processes = {
428            'SynthStrip': self.synthstrip_process,
429            'Coregistration': self.coregistration_process,
430            'Reorientation': self.reorientation_process,
431            'Preprocess': self.dl_preprocess,
432            'Deep Learning': self.dl_process,
433            'Postprocess': self.dl_postprocess
434        }
435
436        for name, process in processes.items():
437            if process is not None and process.state() != QProcess.ProcessState.NotRunning:
438                self.log_update.emit(QCoreApplication.translate("DlWorker", "Stopping {name}...").format(name=name), 'w')
439
440                # No need to disconnect signals manually!
441                process.terminate()
442
443                # only kill if it doesn't respond after a timeout
444                if not process.waitForFinished(2000):
445                    self.log_update.emit(QCoreApplication.translate("DlWorker", "Forcing {name} to quit...").format(name=name), 'e')
446                    process.kill()
447                    process.waitForFinished(1000)
448
449        self.finished.emit(False, QCoreApplication.translate("DlWorker", "Processing cancelled by user"))
450
451    def on_stdout(self, phase, data):
452        """Correct handler for QProcess stdout"""
453        try:
454            # Correctly decode data from QByteArray
455            text = data.data().decode("utf-8", errors='replace').strip()
456            if text:
457                # Split by lines and log each one
458                for line in text.splitlines():
459                    if line.strip():  # Avoid empty lines
460                        self.log_update.emit(f"[{phase}] {line.strip()}", 'i')
461        except Exception as e:
462            self.log_update.emit(f"[{phase}] Error decoding stdout: {str(e)}", 'e')
463
464    def on_stderr(self, phase, data):
465        """Correct handler for QProcess stderr"""
466        try:
467            # Correctly decode data from QByteArray
468            text = data.data().decode("utf-8", errors='replace').strip()
469            if text:
470                # Split by lines and log each one
471                for line in text.splitlines():
472                    if line.strip():  # Avoid empty lines
473                        self.log_update.emit(f"[{phase}] {line.strip()}", 'e')
474        except Exception as e:
475            self.log_update.emit(f"[{phase}] Error decoding stderr: {str(e)}", 'e')
476
477    def on_error(self, phase, error):
478        """Handler for QProcess errors"""
479        try:
480            error_messages = {
481                QProcess.ProcessError.FailedToStart: QCoreApplication.translate("DlWorker", "Failed to start process"),
482                QProcess.ProcessError.Crashed: QCoreApplication.translate("DlWorker", "Process crashed"),
483                QProcess.ProcessError.Timedout: QCoreApplication.translate("DlWorker", "Process timed out"),
484                QProcess.ProcessError.WriteError: QCoreApplication.translate("DlWorker", "Write error"),
485                QProcess.ProcessError.ReadError: QCoreApplication.translate("DlWorker", "Read error"),
486                QProcess.ProcessError.UnknownError: QCoreApplication.translate("DlWorker", "Unknown error")
487            }
488
489            error_msg = error_messages.get(error, QCoreApplication.translate("DlWorker", "Unknown error code: {error}").format(error=error))
490            self.log_update.emit(f"[{phase}] Process error: {error_msg}", 'e')
491
492        except Exception as e:
493            self.log_update.emit(f"[{phase}] Error in error handler: {str(e)}", 'e')