ui.skull_stripping_page

  1import torch
  2from PyQt6.QtWidgets import (
  3    QVBoxLayout, QLabel, QPushButton, QMessageBox, QCheckBox,
  4    QHBoxLayout, QWidget, QDoubleSpinBox, QSpinBox, QGroupBox,
  5    QProgressBar, QFrame
  6)
  7from PyQt6.QtCore import Qt, pyqtSignal, QCoreApplication
  8import os
  9import subprocess
 10
 11from components.file_selector_widget import FileSelectorWidget
 12from threads.skull_strip_thread import SkullStripThread
 13from page import Page
 14from logger import get_logger
 15
 16log = get_logger()
 17
 18
 19class SkullStrippingPage(Page):
 20    """
 21    GUI page for performing skull stripping on NIfTI files.
 22
 23    This class provides a user interface to select anatomical MRI images
 24    and perform skull stripping using either:
 25        - BET (Brain Extraction Tool) from FSL (if available)
 26        - HD-BET (deep learning-based alternative)
 27
 28    It handles:
 29        * User input for tool parameters
 30        * Background processing using a worker thread
 31        * Progress and status updates
 32
 33    Signals:
 34        processing (bool): emitted when processing starts or stops.
 35    """
 36
 37    processing = pyqtSignal(bool)
 38
 39    def __init__(self, context=None, previous_page=None):
 40        """
 41        Initialize the SkullStrippingPage.
 42
 43        Args:
 44            context (dict, optional): Shared app context for state and communication.
 45            previous_page (Page, optional): Reference to previous navigation page.
 46        """
 47        super().__init__()
 48        self.canceled = False
 49        self.context = context
 50        self.previous_page = previous_page
 51        self.next_page = None
 52        self.worker = None  # Background thread reference
 53        self.has_cuda = torch.cuda.is_available()
 54        self.bet_tool = ""
 55
 56        self._setup_ui()
 57        self._translate_ui()
 58
 59        # Re-translate if app language changes
 60        if context and "language_changed" in context:
 61            context["language_changed"].connect(self._translate_ui)
 62
 63    def _setup_ui(self):
 64        """Create and configure all UI components."""
 65        self.processing.connect(self.set_processing_mode)
 66        self.layout = QVBoxLayout(self)
 67        self.setLayout(self.layout)
 68
 69        # Title
 70        self.title = QLabel("Select a NIfTI file for Skull Stripping")
 71        self.title.setStyleSheet("font-size: 18px; font-weight: bold;")
 72        self.title.setAlignment(Qt.AlignmentFlag.AlignCenter)
 73        self.layout.addWidget(self.title)
 74
 75        # Info label (link to BET or HD-BET)
 76        self.info_label = QLabel(
 77            'Using tool: <a href="https://fsl.fmrib.ox.ac.uk/fsl/fslwiki/BET">BET from FSL toolkit</a>'
 78        )
 79        self.info_label.setOpenExternalLinks(True)
 80        self.info_label.setTextInteractionFlags(Qt.TextInteractionFlag.TextBrowserInteraction)
 81        self.info_label.setStyleSheet("""
 82            font-size: 11px;
 83            color: #666;
 84            font-style: italic;
 85            margin-bottom: 6px;
 86        """)
 87        self.info_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
 88        self.info_label.setWordWrap(True)
 89        self.layout.addWidget(self.info_label)
 90
 91        # File selector widget for NIfTI files
 92        self.file_selector_widget = FileSelectorWidget(
 93            parent=self,
 94            context=self.context,
 95            has_existing_function=self.has_existing_skull_strip,
 96            label="skull strip",
 97            allow_multiple=True,
 98            processing=self.processing,
 99            forced_filters={"datatype": "anat"}
100        )
101        self.layout.addWidget(self.file_selector_widget)
102
103        # Check if BET command exists in system
104        try:
105            subprocess.run(['bet'], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
106            self.has_bet = True
107        except FileNotFoundError:
108            self.has_bet = False
109
110        # Setup parameters for BET (if installed)
111        if self.has_bet:
112            self._setup_bet_options()
113        else:
114            try:
115                subprocess.run(['mri_synthstrip'], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
116                self.has_synthstrip = True
117            except FileNotFoundError:
118                self.has_synthstrip = False
119
120            if self.has_synthstrip:
121                self.info_label.setText(
122                    QCoreApplication.translate("SkullStrippingPage", 'Using tool: <a href="https://surfer.nmr.mgh.harvard.edu/docs/synthstrip/">SynthStrip</a> <br>')
123                )
124            else:
125                self.info_label.setText(
126                    QCoreApplication.translate("SkullStrippingPage",
127                        'Using tool: <a href="https://github.com/MIC-DKFZ/HD-BET">hd-bet</a> <br>'
128                        'To use BET from FSL toolkit, install it following the instructions at: '
129                        '<a href="https://fsl.fmrib.ox.ac.uk/fsl/docs/#/install/index">FSL installation</a>')
130                )
131            self.info_label.setOpenExternalLinks(True)
132            self.info_label.setTextInteractionFlags(Qt.TextInteractionFlag.TextBrowserInteraction)
133            self.info_label.setToolTip("Open link")
134
135        # Progress bar
136        self.progress_bar = QProgressBar()
137        self.progress_bar.setVisible(False)
138        self.layout.addWidget(self.progress_bar)
139
140        # Action buttons (Run / Cancel)
141        button_frame = QFrame()
142        button_layout = QHBoxLayout(button_frame)
143        button_layout.setContentsMargins(0, 15, 0, 0)
144
145        # --- RUN button (green) ---
146        self.run_button = QPushButton("Run Skull Stripping")
147        self.file_selector_widget.has_file.connect(self.run_button.setEnabled)
148        self.run_button.setEnabled(False)
149        self.run_button.clicked.connect(self.run_bet)
150        self.run_button.setStyleSheet("""
151                    QPushButton {
152                        font-size: 14px;
153                        font-weight: bold;
154                        background-color: #27ae60;
155                        color: white;
156                        border-radius: 8px;
157                        padding: 10px 20px;
158                        min-width: 140px;
159                    }
160                    QPushButton:hover { 
161                        background-color: #229954; 
162                    }
163                    QPushButton:disabled {
164                        background-color: #bdc3c7;
165                        color: #7f8c8d;
166                    }
167                """)
168        button_layout.addStretch()
169        button_layout.addWidget(self.run_button)
170
171        # --- CANCEL button (red) ---
172        self.cancel_button = QPushButton("Stop Processing")
173        self.cancel_button.setVisible(False)
174        self.cancel_button.clicked.connect(self.cancel_processing)
175        self.cancel_button.setStyleSheet("""
176                    QPushButton {
177                        font-size: 14px;
178                        font-weight: bold;
179                        background-color: #e74c3c;
180                        color: white;
181                        border-radius: 8px;
182                        padding: 10px 20px;
183                        min-width: 140px;
184                    }
185                    QPushButton:hover { 
186                        background-color: #c0392b; 
187                    }
188                    QPushButton:disabled {
189                        background-color: #bdc3c7;
190                        color: #7f8c8d;
191                    }
192                """)
193        button_layout.addWidget(self.cancel_button)
194        button_layout.addStretch()
195
196        self.layout.addWidget(button_frame)
197
198        # Stato
199        self.status_label = QLabel("")
200        self.status_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
201        self.layout.addWidget(self.status_label)
202
203    def _setup_bet_options(self):
204        """Setup BET-specific parameter inputs and advanced options."""
205        # Fractional intensity threshold (main parameter)
206        self.f_box = QGroupBox()
207        f_layout = QHBoxLayout()
208        self.f_label = QLabel("Fractional intensity threshold, smaller values give larger brain outline estimates")
209        f_layout.addWidget(self.f_label)
210
211        self.f_spinbox = QDoubleSpinBox()
212        self.f_spinbox.setRange(0.0, 1.0)
213        self.f_spinbox.setSingleStep(0.05)
214        self.f_spinbox.setValue(0.50)
215        self.f_spinbox.setDecimals(2)
216        self.f_spinbox.setMinimumWidth(60)
217        self.f_spinbox.setMaximumWidth(80)
218        f_layout.addWidget(self.f_spinbox)
219        f_layout.addStretch()
220        self.f_box.setLayout(f_layout)
221        self.layout.addWidget(self.f_box)
222
223        # Toggle for advanced options
224        self.advanced_btn = QPushButton("Show Advanced Options")
225        self.advanced_btn.setCheckable(True)
226        self.advanced_btn.clicked.connect(self.toggle_advanced)
227        self.layout.addWidget(self.advanced_btn)
228
229        # Advanced parameters box
230        self.is_checked = False
231        self.advanced_box = QGroupBox()
232        self.advanced_layout = QVBoxLayout()
233
234        # Output options
235        self.output_label = QLabel("Advanced options:")
236        self.output_label.setStyleSheet("font-weight: bold; margin-top: 5px;")
237        self.advanced_layout.addWidget(self.output_label)
238
239        self.opt_brain_extracted = QCheckBox("Output brain-extracted image")
240        self.opt_brain_extracted.setChecked(True)
241        self.advanced_layout.addWidget(self.opt_brain_extracted)
242
243        self.opt_m = QCheckBox("Output binary brain mask image")
244        self.advanced_layout.addWidget(self.opt_m)
245
246        self.opt_t = QCheckBox("Apply thresholding to brain and mask image")
247        self.advanced_layout.addWidget(self.opt_t)
248
249        self.opt_s = QCheckBox("Output exterior skull surface image")
250        self.advanced_layout.addWidget(self.opt_s)
251
252        self.opt_o = QCheckBox("Output brain surface overlaid onto original image")
253        self.advanced_layout.addWidget(self.opt_o)
254
255        # Threshold gradient
256        self.threshold_layout = QHBoxLayout()
257        self.threshold_label = QLabel("Threshold gradient; positive values give larger brain outline at bottom, smaller at top")
258        self.threshold_layout.addWidget(self.threshold_label)
259        self.g_spinbox = QDoubleSpinBox()
260        self.g_spinbox.setRange(-1.0, 1.0)
261        self.g_spinbox.setSingleStep(0.1)
262        self.g_spinbox.setValue(0.0)
263        self.g_spinbox.setDecimals(1)
264        self.g_spinbox.setMinimumWidth(60)
265        self.g_spinbox.setMaximumWidth(80)
266        self.threshold_layout.addWidget(self.g_spinbox)
267        self.threshold_layout.addStretch()
268        self.advanced_layout.addLayout(self.threshold_layout)
269
270        # Sphere coordinates
271        coords_layout = QHBoxLayout()
272        self.coords_label = QLabel("Coordinates (voxels) for centre of initial brain surface sphere")
273        coords_layout.addWidget(self.coords_label)
274
275        self.c_x_spinbox = QSpinBox()
276        self.c_x_spinbox.setRange(0, 9999)
277        coords_layout.addWidget(self.c_x_spinbox)
278        coords_layout.addWidget(QLabel("Y"))
279
280        self.c_y_spinbox = QSpinBox()
281        self.c_y_spinbox.setRange(0, 9999)
282        coords_layout.addWidget(self.c_y_spinbox)
283        coords_layout.addWidget(QLabel("Z"))
284
285        self.c_z_spinbox = QSpinBox()
286        self.c_z_spinbox.setRange(0, 9999)
287        coords_layout.addWidget(self.c_z_spinbox)
288        coords_layout.addStretch()
289
290        self.advanced_layout.addLayout(coords_layout)
291        self.advanced_box.setLayout(self.advanced_layout)
292        self.advanced_box.setVisible(False)
293        self.layout.addWidget(self.advanced_box)
294
295    def has_existing_skull_strip(self, nifti_file_path, workspace_path):
296        """
297        Check if a skull-stripped file already exists for the selected subject.
298
299        Args:
300            nifti_file_path (str): Full path to input NIfTI file.
301            workspace_path (str): Base workspace directory path.
302
303        Returns:
304            bool: True if an existing skull strip is found, False otherwise.
305        """
306        path_parts = nifti_file_path.replace(workspace_path, '').strip(os.sep).split(os.sep)
307        subject_id = next((p for p in path_parts if p.startswith('sub-')), None)
308        if not subject_id:
309            return False
310
311        skull_strip_dir = os.path.join(workspace_path, 'derivatives', 'skullstrips', subject_id, 'anat')
312        if not os.path.exists(skull_strip_dir):
313            return False
314
315        return any(f.endswith('.nii.gz') for f in os.listdir(skull_strip_dir))
316
317    def toggle_advanced(self):
318        """Toggle visibility of advanced BET parameters."""
319        self.is_checked = self.advanced_btn.isChecked()
320        self.advanced_box.setVisible(self.is_checked)
321        self.advanced_btn.setText("Hide Advanced Options" if self.is_checked else "Show Advanced Options")
322
323    def run_bet(self):
324        """Start the skull stripping process in a background thread."""
325        selected_files = self.file_selector_widget.get_selected_files()
326        if not selected_files:
327            QMessageBox.warning(self, "No files", "Please select at least one NIfTI file first.")
328            return
329
330        parameters = None
331        if self.has_bet:
332            parameters = {
333                'f_val': self.f_spinbox.value(),
334                'opt_brain_extracted': self.opt_brain_extracted.isChecked(),
335                'opt_m': self.opt_m.isChecked(),
336                'opt_t': self.opt_t.isChecked(),
337                'opt_s': self.opt_s.isChecked(),
338                'opt_o': self.opt_o.isChecked(),
339                'g_val': self.g_spinbox.value(),
340                'c_x': self.c_x_spinbox.value(),
341                'c_y': self.c_y_spinbox.value(),
342                'c_z': self.c_z_spinbox.value(),
343            }
344
345        if self.has_bet:
346            self.bet_tool = "fsl-bet"
347        elif self.has_synthstrip:
348            self.bet_tool = "synthstrip"
349        else:
350            self.bet_tool = "hd-bet"
351
352        # Create worker thread
353        self.worker = SkullStripThread(selected_files, self.context["workspace_path"], parameters, self.has_cuda, self.bet_tool)
354
355        # Connect worker signals
356        self.worker.progress_updated.connect(self.on_progress_updated)
357        self.worker.progress_value_updated.connect(self.on_progress_value_updated)
358        self.worker.file_started.connect(self.on_file_started)
359        self.worker.file_completed.connect(self.on_file_completed)
360        self.worker.all_completed.connect(self.on_all_completed)
361        self.worker.finished.connect(self.on_worker_finished)
362
363        # Update UI
364        self.processing.emit(True)
365        self.progress_bar.setRange(0, 100)
366        self.progress_bar.setValue(0)
367        self.progress_bar.setVisible(True)
368
369        self.worker.start()
370
371    def cancel_processing(self):
372        """Cancel the currently running process."""
373        self.canceled = True
374        if self.worker and self.worker.isRunning():
375
376            self.worker.cancel()
377            self.status_label.setText("Cancelling...")
378            self.status_label.setStyleSheet("color: #FF9800; font-weight: bold;")
379
380    def set_processing_mode(self, processing):
381        """Enable or disable UI controls depending on processing state."""
382        self.run_button.setVisible(not processing and bool(self.file_selector_widget.get_selected_files()))
383        self.cancel_button.setVisible(processing)
384        if hasattr(self, "f_spinbox"):
385            self.f_spinbox.setEnabled(not processing)
386        if hasattr(self, "advanced_btn"):
387            self.advanced_btn.setEnabled(not processing)
388        if hasattr(self, "advanced_box"):
389            for widget in self.advanced_box.findChildren(QWidget):
390                widget.setEnabled(not processing)
391        self.cancel_button.setVisible(processing)
392
393
394    def on_progress_updated(self, message):
395        """Update progress text message."""
396        self.status_label.setText(message)
397        self.status_label.setStyleSheet("color: #2196F3; font-weight: bold;")
398
399    def on_progress_value_updated(self, value):
400        """Update numeric progress bar value."""
401        self.progress_bar.setValue(value)
402
403    def on_file_started(self, filename):
404        """Triggered when a new file starts processing."""
405        pass
406
407    def on_file_completed(self, filename, success, error_message):
408        """Triggered when a file finishes processing."""
409        if not success and error_message:
410            log.error(f"Error processing {filename}: {error_message}")
411
412    def on_all_completed(self, success_count, failed_files):
413        """Triggered after all files are processed."""
414        # Nascondi progress bar
415        self.progress_bar.setVisible(False)
416
417        # Aggiorna il messaggio di stato finale
418        if success_count > 0:
419            summary = QCoreApplication.translate("SkullStrippingPage",
420                                                 "Skull Stripping completed successfully for {0} file(s)").format(
421                success_count)
422            if failed_files:
423                failed_summary = QCoreApplication.translate("SkullStrippingPage",
424                                                            "{count} file(s) failed: {files}").format(
425                    count=len(failed_files),
426                    files=', '.join([os.path.basename(f) for f in failed_files])
427                )
428                summary += f"\n{failed_summary}"
429                self.status_label.setStyleSheet("color: #FF9800; font-weight: bold; font-size: 12pt; padding: 5px;")
430            else:
431                self.status_label.setStyleSheet("color: #4CAF50; font-weight: bold; font-size: 12pt; padding: 5px;")
432        else:
433            failed_files_num = len(failed_files)
434            if failed_files_num == 0:
435                failed_files_num = ""
436            summary = QCoreApplication.translate("SkullStrippingPage", "All {0} file(s) failed to process").format(
437                failed_files_num)
438            self.status_label.setStyleSheet("color: #FF0000; font-weight: bold; font-size: 12pt; padding: 5px;")
439
440        self.status_label.setText(summary)
441
442
443    def on_worker_finished(self):
444            """Triggered when worker thread finishes execution."""
445            self.processing.emit(False)
446            if self.worker:
447                self.worker.deleteLater()
448                self.worker = None
449            if self.context and "update_main_buttons" in self.context:
450                self.context["update_main_buttons"]()
451
452    def back(self):
453        """Navigate back to the previous page if not processing."""
454        if self.worker and self.worker.isRunning():
455            QMessageBox.warning(self, "Processing in progress",
456                                "Cannot go back while skull stripping is in progress. Please wait or cancel.")
457            log.warning("Processing in progress")
458            return None
459        if self.previous_page:
460            self.previous_page.on_enter()
461            return self.previous_page
462        return None
463
464    def on_enter(self):
465        """Reset status label when page is entered."""
466        self.status_label.setText("")
467
468    def is_ready_to_advance(self):
469        """Return False since this page doesn't advance automatically."""
470        return False
471
472    def is_ready_to_go_back(self):
473        """Disable back navigation during processing."""
474        return not (self.worker and self.worker.isRunning())
475
476    def reset_page(self):
477        """Reset the page state to its initial configuration."""
478        if self.worker and self.worker.isRunning():
479            self.worker.cancel()
480            self.worker.wait()
481
482        self.file_selector_widget.clear_selected_files()
483        self.run_button.setEnabled(False)
484
485        # Reset advanced options
486        self.f_spinbox.setValue(0.50)
487        self.advanced_btn.setChecked(False)
488        self.advanced_box.setVisible(False)
489        self.advanced_btn.setText("Show Advanced Options")
490
491        self.opt_brain_extracted.setChecked(True)
492        self.opt_m.setChecked(False)
493        self.opt_t.setChecked(False)
494        self.opt_s.setChecked(False)
495        self.opt_o.setChecked(False)
496
497        self.g_spinbox.setValue(0.0)
498        self.c_x_spinbox.setValue(0)
499        self.c_y_spinbox.setValue(0)
500        self.c_z_spinbox.setValue(0)
501
502        self.progress_bar.setVisible(False)
503        self.cancel_button.setVisible(False)
504        self.processing.emit(False)
505        self.status_label.setText("")
506
507    def _translate_ui(self):
508        """Translate UI strings according to the active language."""
509        self.title.setText(QCoreApplication.translate("SkullStrippingPage", "Select a NIfTI file for Skull Stripping"))
510        if self.has_bet:
511            self.f_label.setText(QCoreApplication.translate("SkullStrippingPage", "Fractional intensity threshold, smaller values give larger brain outline estimates"))
512            self.advanced_btn.setText(QCoreApplication.translate("SkullStrippingPage", "Show Advanced Options"))
513            self.output_label.setText(QCoreApplication.translate("SkullStrippingPage", "Advanced options:"))
514            self.opt_brain_extracted.setText(QCoreApplication.translate("SkullStrippingPage", "Output brain-extracted image"))
515            self.opt_m.setText(QCoreApplication.translate("SkullStrippingPage", "Output binary brain mask image"))
516            self.opt_t.setText(QCoreApplication.translate("SkullStrippingPage", "Apply thresholding to brain and mask image"))
517            self.opt_s.setText(QCoreApplication.translate("SkullStrippingPage", "Output exterior skull surface image"))
518            self.opt_o.setText(QCoreApplication.translate("SkullStrippingPage", "Output brain surface overlaid onto original image"))
519            self.threshold_label.setText(QCoreApplication.translate("SkullStrippingPage", "Threshold gradient; positive values give larger brain outline at bottom, smaller at top"))
520            self.coords_label.setText(QCoreApplication.translate("SkullStrippingPage", "Coordinates (voxels) for centre of initial brain surface sphere"))
521        else:
522            self.info_label.setToolTip(QCoreApplication.translate("SkullStrippingPage", "Open link"))
523
524        self.run_button.setText(QCoreApplication.translate("SkullStrippingPage", "Run Skull Stripping"))
525        self.cancel_button.setText(QCoreApplication.translate("SkullStrippingPage", "Cancel"))