ui.dl_execution_page

  1import os
  2import subprocess
  3
  4import torch
  5from PyQt6 import QtWidgets, QtGui, QtCore
  6from PyQt6.QtGui import QFont
  7from PyQt6.QtWidgets import (QWidget, QVBoxLayout, QLabel, QPushButton,
  8                             QScrollArea, QFrame, QGridLayout, QHBoxLayout,
  9                             QMessageBox, QGroupBox, QListWidget, QProgressBar,
 10                             QListWidgetItem, QTextEdit, QSplitter, QFileDialog,
 11                             QCheckBox)
 12from PyQt6.QtCore import Qt, QThread, QCoreApplication
 13
 14from components.circular_progress_bar import CircularProgress
 15from threads.dl_worker import DlWorker
 16from page import Page
 17from logger import get_logger
 18
 19log = get_logger()
 20
 21
 22class DlExecutionPage(Page):
 23    """Page for SynthStrip + Coregistration of selected NIfTI files."""
 24
 25    def __init__(self, context=None, previous_page=None):
 26        super().__init__()
 27        self.worker = None  # Worker thread instance
 28        self.current_file = None  # Currently processing file (if needed)
 29        self.context = context  # Shared application context (settings, file lists)
 30        self.previous_page = previous_page  # Reference to the previous page for navigation
 31        self.next_page = None  # Reference to the next page
 32
 33        # State flags
 34        self.processing = False  # Flag to indicate if processing is active
 35        self.processing_completed = False  # Flag to indicate if processing has finished
 36
 37        self._setup_ui()  # Initialize the user interface elements
 38
 39        self._translate_ui()  # Apply initial translations
 40        if context and "language_changed" in context:
 41            # Connect to the global language change signal if it exists
 42            context["language_changed"].connect(self._translate_ui)
 43
 44    def _setup_ui(self):
 45        """Sets up the user interface layout and widgets."""
 46        main_layout = QVBoxLayout(self)
 47
 48        # --- Header ---
 49        self.header = QLabel("Deep Learning Segmentation")
 50        self.header.setFont(QFont("Arial", 18, QFont.Weight.Bold))
 51        self.header.setAlignment(Qt.AlignmentFlag.AlignCenter)
 52        main_layout.addWidget(self.header)
 53
 54        # --- Info Label (Tool detection) ---
 55        self.info_label = QLabel("")
 56        self.info_label.setOpenExternalLinks(True)
 57        self.info_label.setTextInteractionFlags(Qt.TextInteractionFlag.TextBrowserInteraction)
 58
 59        self.info_label.setStyleSheet("""
 60                            font-size: 11px;
 61                            color: #666;
 62                            font-style: italic;
 63                            margin-bottom: 6px;
 64                        """)
 65        self.info_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
 66        self.info_label.setWordWrap(True)
 67        main_layout.addWidget(self.info_label)
 68
 69        try:
 70            # Run the command with --help to check if it exists and works
 71            subprocess.run(
 72                ['mri_synthstrip', '--help'],
 73                stdout=subprocess.DEVNULL,  # Discard standard output
 74                stderr=subprocess.DEVNULL  # Discard standard error
 75            )
 76            self.has_freesurfer = True
 77            self.info_label.setText(
 78                "Using tool: <a href='https://surfer.nmr.mgh.harvard.edu/docs/synthstrip/'>SynthStrip from Free Surfer</a>")
 79        except (FileNotFoundError, subprocess.CalledProcessError):
 80            # Handle case where FreeSurfer version is not found
 81            self.has_freesurfer = False
 82            self.info_label.setText(
 83                'Using tool: <a href="https://github.com/nipreps/synthstrip">SynthStrip (nipreps)</a><br>'
 84                'To use SynthStrip from FreeSurfer, follow the instructions at: '
 85                '<a href="https://surfer.nmr.mgh.harvard.edu/docs/synthstrip/">SynthStrip official documentation</a>'
 86            )
 87
 88        self.info_label.setOpenExternalLinks(True)  # Enable opening links in browser
 89        self.info_label.setTextInteractionFlags(Qt.TextInteractionFlag.TextBrowserInteraction)
 90        self.info_label.setToolTip("Open link")
 91
 92        # --- Current Operation Status ---
 93        self.current_operation = QLabel("Ready to start")
 94        self.current_operation.setStyleSheet("""
 95            font-size: 13px;
 96            color: #7f8c8d;
 97            margin-top: 8px;
 98        """)
 99        self.current_operation.setAlignment(Qt.AlignmentFlag.AlignCenter)
100        main_layout.addWidget(self.current_operation)
101
102        # --- Content area (Grid) ---
103        content_layout = QGridLayout()
104        main_layout.addLayout(content_layout, stretch=1)
105
106        # --- Left Column (Row 0, Col 0): Progress Display ---
107        left_layout = QVBoxLayout()
108        self.progress_bar = CircularProgress()  # Custom circular progress widget
109        left_layout.addWidget(self.progress_bar, alignment=Qt.AlignmentFlag.AlignCenter)
110        content_layout.addLayout(left_layout, 0, 0)
111
112        # --- Right Column (Row 0, Col 1): File List ---
113        self.files_group = QGroupBox("Files to process")
114        files_layout = QVBoxLayout(self.files_group)
115
116        self.files_list = QListWidget()
117        self.files_list.setMaximumHeight(150)
118        files_layout.addWidget(self.files_list)
119
120        # Add the file list (QListWidget) directly to the grid
121        content_layout.addWidget(self.files_list, 0, 1)
122
123        # --- Grid Layout Stretch ---
124        # Columns: 1/3 for progress bar, 2/3 for scroll area
125        content_layout.setColumnStretch(0, 1)  # Left column (progress bar)
126        content_layout.setColumnStretch(1, 2)  # Right column (folder list)
127
128        # --- Bottom Row (Rows 1 & 2): Log Display ---
129        self.log_label = QLabel("Execution Log:")
130        self.log_label.setStyleSheet("""
131                    font-size: 16px;
132                    font-weight: bold;
133                    color: #2c3e50;
134                    margin-top: 15px;
135                    margin-bottom: 5px;
136                """)
137        # Add log label, spanning both columns (Row 1, Col 0, 1 row span, 2 col span)
138        content_layout.addWidget(self.log_label, 1, 0, 1, 2)
139
140        self.log_text = QTextEdit()
141        self.log_text.setStyleSheet("""
142                    QTextEdit {
143                        background-color: #2c3e50;
144                        color: #ecf0f1;
145                        font-family: 'Courier New', monospace;
146                        font-size: 11px;
147                        border: 2px solid #34495e;
148                        border-radius: 8px;
149                        padding: 8px;
150                    }
151                """)
152        self.log_text.setMaximumHeight(200)
153        self.log_text.setReadOnly(True)
154        # Add log text area, spanning both columns (Row 2, Col 0, 1 row span, 2 col span)
155        content_layout.addWidget(self.log_text, 2, 0, 1, 2)
156
157        # === Control Buttons ===
158
159        # --- Button Layout (Start/Stop) ---
160        button_frame = QFrame()  # Frame to hold the centered buttons
161        button_layout = QHBoxLayout(button_frame)
162        button_layout.setContentsMargins(0, 15, 0, 0)
163
164        self.start_button = QPushButton("Start deep learning")
165        button_layout.addStretch()  # Center the buttons
166        self.start_button.clicked.connect(self.start_processing)
167        self.start_button.setStyleSheet("""
168            QPushButton {
169                font-size: 14px;
170                font-weight: bold;
171                background-color: #27ae60;
172                color: white;
173                border-radius: 8px;
174                padding: 10px 20px;
175                min-width: 140px;
176            }
177            QPushButton:hover { 
178                background-color: #229954; 
179            }
180            QPushButton:disabled {
181                background-color: #bdc3c7;
182                color: #7f8c8d;
183            }
184        """)
185        button_layout.addWidget(self.start_button)
186
187        self.cancel_button = QPushButton("Stop processing")
188        self.cancel_button.setStyleSheet("""
189            QPushButton {
190                font-size: 14px;
191                font-weight: bold;
192                background-color: #e74c3c;
193                color: white;
194                border-radius: 8px;
195                padding: 10px 20px;
196                min-width: 140px;
197            }
198            QPushButton:hover { 
199                background-color: #c0392b; 
200            }
201            QPushButton:disabled {
202                background-color: #bdc3c7;
203                color: #7f8c8d;
204            }
205        """)
206        self.cancel_button.clicked.connect(self.cancel_processing)
207        self.cancel_button.setVisible(False)  # Initially hidden
208
209        button_layout.addWidget(self.cancel_button)
210        button_layout.addStretch()  # Center the buttons
211
212        main_layout.addWidget(button_frame)
213
214    def on_enter(self):
215        """Called when this page becomes active."""
216        self.reset_processing_state()  # Ensure UI is in a clean, default state
217
218        # Populate the file list if available from context
219        if self.context and "selected_segmentation_files" in self.context:
220            self.files_list.clear()
221            for file_path in self.context["selected_segmentation_files"]:
222                filename = os.path.basename(file_path)
223                # Add item with "Waiting..." status
224                self.files_list.addItem(
225                    QCoreApplication.translate("DlExecutionPage", "📄 {filename} - Waiting...").format(
226                        filename=filename))
227
228    def start_processing(self):
229        """Starts the deep learning processing workflow."""
230        if not self.context or "selected_segmentation_files" not in self.context:
231            QMessageBox.warning(self, QCoreApplication.translate("DlExecutionPage", "Error"),
232                                QCoreApplication.translate("DlExecutionPage", "No files selected for processing."))
233            return
234
235        selected_files = self.context["selected_segmentation_files"]
236        if not selected_files:
237            QMessageBox.warning(self, QCoreApplication.translate("DlExecutionPage", "Error"),
238                                QCoreApplication.translate("DlExecutionPage", "No files selected for processing."))
239            return
240
241        # --- Setup Worker Thread ---
242        self.worker = DlWorker(
243            input_files=selected_files,
244            workspace_path=self.context["workspace_path"],
245            has_freesurfer=self.has_freesurfer
246        )
247
248        # --- Connect Worker Signals ---
249        self.worker.progressbar_update.connect(self.update_progress)
250        self.worker.file_update.connect(self.update_file_status)
251        self.worker.log_update.connect(self.add_log_message)
252        self.worker.finished.connect(self.processing_finished)
253        # self.worker.finished.connect(self.worker.deleteLater) # Option to auto-delete worker
254
255        # --- Update UI for Processing State ---
256        self.processing = True
257        self.start_button.setVisible(False)
258        self.cancel_button.setVisible(True)
259        self.progress_bar.setVisible(True)
260        self.progress_bar.setValue(0)
261        self.current_operation.setText(QCoreApplication.translate("DlExecutionPage", "Processing..."))
262        self.log_text.clear()
263
264        # --- Start Worker ---
265        self.worker.start()  # Start the QThread
266
267        self.add_log_message(
268            QCoreApplication.translate("DlExecutionPage", "Deep learning processing started for {0} file").format(
269                len(selected_files)), 'i')
270        log.info(f"Deep learning processing started for {len(selected_files)} file")
271
272    def update_progress(self, value):
273        """Updates the circular progress bar value."""
274        self.progress_bar.setValue(value)
275
276    def add_log_message(self, message, type):
277        """Appends a message to the UI log text area."""
278        timestamp = QtCore.QDateTime.currentDateTime().toString("hh:mm:ss")
279        self.log_text.append(f"[{timestamp}] {message}")
280
281        # Log to external file based on message type
282        if type == 'e':
283            log.error(f"[{timestamp}] {message}")
284        elif type == 'i':
285            log.info(f"[{timestamp}] {message}")
286        elif type == 'w':
287            log.warning(f"[{timestamp}] {message}")
288        else:
289            log.debug(f"[{timestamp}] {message}")
290
291        # Auto-scroll to the bottom
292        scrollbar = self.log_text.verticalScrollBar()
293        scrollbar.setValue(scrollbar.maximum())
294
295    def update_file_status(self, filename, status):
296        """Updates the status of a specific file in the QListWidget."""
297        for i in range(self.files_list.count()):
298            item = self.files_list.item(i)
299            if filename in item.text():  # Find the matching item
300                item.setText(f"📄 {filename} - {status}")
301                break  # Stop searching once found
302
303    def cancel_processing(self):
304        """Requests the worker thread to cancel processing."""
305        if self.worker:
306            self.worker.cancel_requested.emit()  # Emit the signal to request cancellation
307            # self.worker.cancel() # If you have a hard cancel method
308            self.add_log_message(QCoreApplication.translate("DlExecutionPage", "Cancellation requested..."), 'i')
309
310    def processing_finished(self, success, message):
311        """Slot called when the worker thread finishes."""
312        self.processing = False
313        self.processing_completed = True
314
315        # --- Update UI to Reflect Finished State ---
316        self.start_button.setVisible(True)
317        self.cancel_button.setVisible(False)
318        self.progress_bar.setVisible(False)
319
320        if success:
321            self.current_operation.setText(QCoreApplication.translate("DlExecutionPage", "✓ Processing completed!"))
322            self.current_operation.setStyleSheet("color: green; font-weight: bold;")
323            self.start_button.setText(QCoreApplication.translate("DlExecutionPage", "Reprocess"))
324        else:
325            self.current_operation.setText("✗ Processing failed!")
326            self.current_operation.setStyleSheet("color: #c0392b; font-weight: bold;")
327
328        self.add_log_message(QCoreApplication.translate("DlExecutionPage", "Final: {message}").format(message=message),
329                             'i')
330
331        # --- Update Context with Results ---
332        if success and "workspace_path" in self.context:
333            # If successful, store the output directory in the context
334            output_dir = os.path.join(self.context["workspace_path"], "outputs")
335            self.context["processing_output_dir"] = output_dir
336
337        # Notify main window to update its buttons (e.g., 'Next')
338        if "update_main_buttons" in self.context:
339            self.context["update_main_buttons"]()
340
341        # --- Show Final Message Box ---
342        if success:
343            QMessageBox.information(self, QCoreApplication.translate("DlExecutionPage", "Completed"), message)
344        else:
345            QMessageBox.critical(self, QCoreApplication.translate("DlExecutionPage", "Error"), message)
346
347    def reset_processing_state(self):
348        """Resets the page to its initial state."""
349        self.processing = False
350        self.processing_completed = False
351        self.start_button.setText(QCoreApplication.translate("DlExecutionPage", "Start deep learning"))
352        self.start_button.setVisible(True)
353        self.cancel_button.setVisible(False)
354        self.progress_bar.setVisible(False)
355        self.current_operation.setStyleSheet("")  # Reset color
356        self.current_operation.setText(QCoreApplication.translate("DlExecutionPage", "Ready to start"))
357
358        if self.worker:
359            self.worker = None  # Clear worker reference
360
361    def back(self):
362        """Handles the 'Back' navigation request."""
363        if self.processing:
364            # Ask user for confirmation to stop
365            reply = QMessageBox.question(
366                self,
367                QCoreApplication.translate("DlExecutionPage", "Processing in progress"),
368                QCoreApplication.translate("DlExecutionPage",
369                                           "Processing is in progress. Do you really want to go back?\nProcessing will be interrupted."),
370                QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No,
371                QMessageBox.StandardButton.No
372            )
373            if reply == QMessageBox.StandardButton.No:
374                return None  # Abort navigation
375
376            # --- Stop Processing ---
377            if self.worker:
378                self.worker.cancel_requested.emit()
379                # self.worker.cancel()
380
381        if self.previous_page:
382            self.previous_page.on_enter()  # Call on_enter of the previous page
383            return self.previous_page
384        return None  # No previous page
385
386    def next(self, context):
387        """Handles the 'Next' navigation request."""
388        return None  # This page is a terminal step
389
390    def is_ready_to_advance(self):
391        """Checks if the page is in a state to advance."""
392        return False  # Cannot advance from this page
393
394    def is_ready_to_go_back(self):
395        """Checks if the page is in a state to go back."""
396        return True  # Can always go back (with a warning if processing)
397
398    def _translate_ui(self):
399        """Updates all user-visible strings for translation."""
400        self.header.setText(QCoreApplication.translate("DlExecutionPage", "Deep Learning Segmentation"))
401
402        # Only reset 'Ready to start' if not processing or completed
403        if not self.processing and not self.processing_completed:
404            self.current_operation.setText(QCoreApplication.translate("DlExecutionPage", "Ready to start"))
405
406        self.files_group.setTitle(QCoreApplication.translate("DlExecutionPage", "Files to process"))
407        self.log_label.setText(QCoreApplication.translate("DlExecutionPage", "Execution Log:"))
408
409        if not self.processing_completed:
410            self.start_button.setText(QCoreApplication.translate("DlExecutionPage", "Start deep learning"))
411        else:
412            self.start_button.setText(QCoreApplication.translate("DlExecutionPage", "Reprocess"))
413
414        self.cancel_button.setText(QCoreApplication.translate("DlExecutionPage", "Stop processing"))
415
416        # Repopulate the file list with translated "Waiting..." string
417        if self.context and "selected_segmentation_files" in self.context and not self.processing:
418            self.files_list.clear()
419            for file_path in self.context["selected_segmentation_files"]:
420                filename = os.path.basename(file_path)
421                self.files_list.addItem(
422                    QCoreApplication.translate("DlExecutionPage", "📄 {filename} - Waiting...").format(
423                        filename=filename))