Coverage for main/threads/dl_worker.py: 55%
287 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-03 15:51 +0000
1import os
2import shutil
3import sys
4import tempfile
5from pathlib import Path
7from PyQt6.QtCore import QThread, pyqtSignal, QProcess, QObject, QCoreApplication
9from logger import get_logger
10from utils import get_bin_path, resource_path, get_dl_python_executable, get_script_path
12log = get_logger()
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()
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
28 self.output_dir = None
30 self.total_files = None
31 self.processed_files = None
32 self.failed_files = None
34 self.is_cancelled = False
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
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")
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
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}")
63 self.python_executable = get_dl_python_executable()
65 self.cancel_requested.connect(self.cancel)
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
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
76 progress = int((total_phases_completed / total_phases_overall) * 100)
77 progress = min(progress, 100) # Ensure it doesn't exceed 100
79 self.progressbar_update.emit(progress)
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 )
88 def start(self):
89 """Processes all NIfTI files with the Deep Learning pipeline"""
91 self.total_files = len(self.input_files)
92 log.debug(self.input_files)
93 self.processed_files = 0
94 self.failed_files = []
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)
102 self.process_single_file()
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
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)
114 self.log_update.emit(QCoreApplication.translate("DlWorker", "=== PROCESSING: {0} ===").format(self.current_input_file_basename), 'i')
116 # PHASE 1: SynthStrip
117 self.current_phase = 1
118 self.run_synthstrip()
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')
126 self.log_update.emit(QCoreApplication.translate("DlWorker", "SynthStrip started: {0}").format(os.path.basename(self.current_input_file)), 'i')
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")
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()))
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 ]
153 self.synthstrip_process.start(cmd[0], cmd[1:])
155 def on_synthstrip_finished(self, exit_code, exit_status):
156 if self.is_cancelled:
157 return
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
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
174 def run_coregistration(self):
175 """Performs coregistration with atlas"""
176 phase = "Coregistration"
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')
181 self.log_update.emit(QCoreApplication.translate("DlWorker", "Coregistration started: {0}").format(os.path.basename(self.current_input_file)), 'i')
183 # Create directory for coregistration
184 coreg_dir = os.path.join(self.output_dir, "coregistration")
185 os.makedirs(coreg_dir, exist_ok=True)
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()))
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 ]
201 self.coregistration_process.start(self.python_executable, args)
203 def on_coregistration_finished(self, exit_code, exit_status):
204 if self.is_cancelled:
205 return
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
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
222 def run_reorientation(self):
223 """Performs reorientation of the brain_in_atlas file using the BraTS affine matrix"""
224 phase = "Reorientation"
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')
229 self.log_update.emit(QCoreApplication.translate("DlWorker", "Reorientation started: {0}").format(os.path.basename(self.current_input_file)), 'i')
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"))
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
238 rsl_files = [f for f in brain_in_atlas_files if f.name.endswith("_rsl.nii") or f.name.endswith("_rsl.nii.gz")]
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])
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()))
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 ]
260 self.reorientation_process.start(self.python_executable, args)
262 def on_reorientation_finished(self, exit_code, exit_status):
263 if self.is_cancelled:
264 return
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
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
281 def run_preprocess(self):
282 """Runs PHASE 4: PREPARE and PREPROCESS"""
283 phase = "Preprocessing"
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')
288 data_path = os.path.join(self.output_dir, "reoriented")
289 results_path = os.path.join(self.output_dir, "preprocess")
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()))
297 args = [
298 get_script_path("deep_learning/preprocess.py"),
299 '--data', data_path,
300 '--results', results_path,
301 '--ohe'
302 ]
304 self.dl_preprocess.start(self.python_executable, args)
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
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
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
325 def run_deep_learning(self):
326 """Runs PHASE 5: DEEP LEARNING"""
327 phase = "Deep Learning execution"
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')
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()))
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 ]
353 self.dl_process.start(self.python_executable, args)
355 def on_dl_finished(self, exit_code, exit_status):
356 if self.is_cancelled:
357 return
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
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
374 def run_postprocess(self):
375 phase = "Postprocessing"
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')
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()))
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 ]
396 self.dl_postprocess.start(self.python_executable, args)
398 def on_postprocess_finished(self, exit_code, exit_status):
399 if self.is_cancelled:
400 return
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
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()
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
423 def cancel(self):
424 self.is_cancelled = True
425 self.log_update.emit(QCoreApplication.translate("DlWorker", "Cancellation requested - stopping all processes..."), 'w')
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 }
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')
440 # No need to disconnect signals manually!
441 process.terminate()
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)
449 self.finished.emit(False, QCoreApplication.translate("DlWorker", "Processing cancelled by user"))
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')
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')
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 }
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')
492 except Exception as e:
493 self.log_update.emit(f"[{phase}] Error in error handler: {str(e)}", 'e')