components.nifti_file_dialog

  1from pathlib import Path
  2
  3from PyQt6.QtWidgets import (
  4    QDialog, QVBoxLayout, QGroupBox, QGridLayout, QLabel, QLineEdit, QComboBox,
  5    QCheckBox, QListWidget, QListWidgetItem, QPushButton, QDialogButtonBox,
  6    QMessageBox
  7)
  8from PyQt6.QtCore import Qt, QCoreApplication
  9from PyQt6.QtGui import QBrush, QColor
 10import os
 11from logger import get_logger
 12
 13log = get_logger()
 14
 15class NiftiFileDialog(QDialog):
 16    def __init__(self, context, allow_multiple=None, has_existing_func=None, label=None, forced_filters=None):
 17        super().__init__()
 18        self.label = label
 19        self.setWindowTitle(QCoreApplication.translate("Components", "Select NIfTI Files ({0})").format(self.label))
 20        self.resize(700, 650)
 21
 22        self.workspace_path = context["workspace_path"]
 23        self.allow_multiple = allow_multiple
 24        self.has_existing_func = has_existing_func or (lambda *_: False)
 25        self.forced_filters = forced_filters or {}
 26
 27        self.selected_files = []
 28        self.relative_to_absolute = {}
 29        self.files_with_flag = set()
 30
 31        self._build_ui()
 32        self._populate_files()
 33        self._connect_signals()
 34
 35    # === Public API ===
 36    @staticmethod
 37    def get_files(context, allow_multiple=True, has_existing_func=None, label="mask",forced_filters=None):
 38        dialog = NiftiFileDialog(context, allow_multiple, has_existing_func, label,forced_filters)
 39        if dialog.exec() == QDialog.DialogCode.Accepted:
 40            return dialog.selected_files
 41        return None
 42
 43    # === UI Construction ===
 44    def _build_ui(self):
 45        layout = QVBoxLayout(self)
 46
 47        # === Filters ===
 48        filter_group = QGroupBox(QCoreApplication.translate("Components", "Filters"))
 49        filter_layout = QGridLayout(filter_group)
 50
 51        self.search_bar = QLineEdit()
 52        self.search_bar.setPlaceholderText(QCoreApplication.translate("Components", "Search files (FLAIR, T1, T2, etc.)"))
 53        self.search_bar.setClearButtonEnabled(True)
 54        filter_layout.addWidget(QLabel(QCoreApplication.translate("Components", "Search:")), 0, 0)
 55        filter_layout.addWidget(self.search_bar, 0, 1, 1, 3)
 56
 57        self.subject_combo = QComboBox()
 58        self.subject_combo.addItem(QCoreApplication.translate("Components", "All subjects"))
 59        self.subject_combo.model().item(0)
 60        filter_layout.addWidget(QLabel(QCoreApplication.translate("Components", "Subject:")), 1, 0)
 61        filter_layout.addWidget(self.subject_combo, 1, 1)
 62
 63        self.session_combo = QComboBox()
 64        self.session_combo.addItem(QCoreApplication.translate("Components", "All sessions"))
 65        self.session_combo.model().item(0)
 66        filter_layout.addWidget(QLabel(QCoreApplication.translate("Components", "Session:")), 1, 2)
 67        filter_layout.addWidget(self.session_combo, 1, 3)
 68
 69        self.modality_combo = QComboBox()
 70        self.modality_combo.addItem(QCoreApplication.translate("Components", "All modalities"))
 71        self.modality_combo.model().item(0)
 72        filter_layout.addWidget(QLabel(QCoreApplication.translate("Components", "Modality:")), 2, 0)
 73        filter_layout.addWidget(self.modality_combo, 2, 1)
 74
 75        self.datatype_combo = QComboBox()
 76        self.datatype_combo.addItem(QCoreApplication.translate("Components", "All types"))
 77        self.datatype_combo.model().item(0)
 78        filter_layout.addWidget(QLabel(QCoreApplication.translate("Components", "Data type:")), 2, 2)
 79        filter_layout.addWidget(self.datatype_combo, 2, 3)
 80
 81        self.no_flag_checkbox = QCheckBox(QCoreApplication.translate("Components", "Show only files without existing {0}s").format(self.label))
 82        self.with_flag_checkbox = QCheckBox(QCoreApplication.translate("Components", "Show only files with existing {0}s").format(self.label))
 83        if self.label is None:
 84            self.no_flag_checkbox.setVisible(False)
 85            self.with_flag_checkbox.setVisible(False)
 86        filter_layout.addWidget(self.no_flag_checkbox, 3, 0, 1, 2)
 87        filter_layout.addWidget(self.with_flag_checkbox, 3, 2, 1, 2)
 88
 89        layout.addWidget(filter_group)
 90
 91        # === Info label ===
 92        self.info_label = QLabel()
 93        layout.addWidget(self.info_label)
 94
 95        # === File list ===
 96        self.file_list = QListWidget()
 97        self.file_list.setEditTriggers(QListWidget.EditTrigger.NoEditTriggers)
 98        self.file_list.setSelectionMode(
 99            QListWidget.SelectionMode.ExtendedSelection if self.allow_multiple else QListWidget.SelectionMode.SingleSelection
100        )
101        self.file_list.setAlternatingRowColors(True)
102        layout.addWidget(self.file_list)
103
104        # === Buttons ===
105        self.buttons = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel)
106
107        reset_button = QPushButton(QCoreApplication.translate("Components", "Reset Filters"))
108        select_all_button = QPushButton(QCoreApplication.translate("Components", "Select All Visible"))
109        deselect_all_button = QPushButton(QCoreApplication.translate("Components", "Deselect All"))
110
111        self.buttons.addButton(reset_button, QDialogButtonBox.ButtonRole.ResetRole)
112        self.buttons.addButton(select_all_button, QDialogButtonBox.ButtonRole.ActionRole)
113        self.buttons.addButton(deselect_all_button, QDialogButtonBox.ButtonRole.ActionRole)
114
115        layout.addWidget(self.buttons)
116
117        # === Store buttons for handlers ===
118        self._reset_button = reset_button
119        self._select_all_button = select_all_button
120        self._deselect_all_button = deselect_all_button
121
122    # === Populate ===
123    def _populate_files(self):
124        self.all_nii_files = []
125        subjects_set, sessions_set, modalities_set, datatypes_set = set(), set(), set(), set()
126
127        for root, dirs, files in os.walk(self.workspace_path):
128            # dirs[:] = [d for d in dirs if d != "derivatives"]
129
130            for f in files:
131                if f.endswith((".nii", ".nii.gz")):
132                    full_path = os.path.join(root, f)
133                    relative_path = os.path.relpath(full_path, self.workspace_path)
134                    self.all_nii_files.append(relative_path)
135                    self.relative_to_absolute[relative_path] = full_path
136
137                    path_parts = relative_path.split(os.sep)
138
139                    # subject
140                    for part in path_parts:
141                        if part.startswith("sub-"):
142                            subjects_set.add(part)
143                            break
144
145                    # session
146                    for part in path_parts:
147                        if part.startswith("ses-"):
148                            sessions_set.add(part)
149                            break
150
151                    # modality
152                    filename_noext = os.path.splitext(os.path.splitext(f)[0])[0]
153                    modality = filename_noext.split("_")[-1] or QCoreApplication.translate("Components", "Unknown")
154                    modalities_set.add(modality)
155
156                    # existing flag (mask/skull strip)
157                    if self.has_existing_func(full_path, self.workspace_path):
158                        self.files_with_flag.add(relative_path)
159
160        # --- Cerca data type directories ---
161        for subject_dir in sorted(Path(self.workspace_path).glob("sub-*")):
162            if subject_dir.is_dir():
163                # Caso: sub-*/anat, sub-*/func, ecc.
164                for child in subject_dir.iterdir():
165                    if child.is_dir() and not child.name.startswith("ses-"):
166                        datatypes_set.add(child.name)
167
168                # Caso: sub-*/ses-*/anat, sub-*/ses-*/func, ecc.
169                for ses_dir in subject_dir.glob("ses-*"):
170                    if ses_dir.is_dir():
171                        for child in ses_dir.iterdir():
172                            if child.is_dir():
173                                datatypes_set.add(child.name)
174
175        self.subject_combo.addItems(sorted(subjects_set))
176        self.session_combo.addItems(sorted(sessions_set))
177        self.modality_combo.addItems(sorted(modalities_set))
178        self.datatype_combo.addItems(sorted(datatypes_set))
179
180        self._update_info_label(len(self.all_nii_files))
181        self._populate_file_list()
182
183        if self.forced_filters:
184            self._apply_forced_filters()
185
186    def _populate_file_list(self):
187        self.file_list.clear()
188        for relative_path in sorted(self.all_nii_files):
189            item = QListWidgetItem(relative_path)
190            if relative_path in self.files_with_flag:
191                item.setForeground(QBrush(QColor(255, 193, 7)))
192                item.setToolTip(QCoreApplication.translate("Components", "{0}\n✓ Existing {1}").format(self.relative_to_absolute[relative_path], self.label))
193            else:
194                item.setToolTip(QCoreApplication.translate("Components", "{0}\n○ No {1}").format(self.relative_to_absolute[relative_path], self.label))
195            self.file_list.addItem(item)
196
197    # === Helpers ===
198    def _update_info_label(self, visible_count):
199        info_text = QCoreApplication.translate("Components", "Showing {0} of {1} files").format(visible_count, len(self.all_nii_files))
200        if self.files_with_flag:
201            info_text += QCoreApplication.translate("Components", " ({0} with existing {1}s)").format(len(self.files_with_flag), self.label)
202        self.info_label.setText(info_text)
203        self.info_label.setStyleSheet("color: gray; font-size: 10px;")
204
205    def _apply_filters(self):
206        search_text = self.search_bar.text().lower()
207        selected_subject = self.subject_combo.currentText()
208        selected_session = self.session_combo.currentText()
209        selected_modality = self.modality_combo.currentText()
210        selected_datatype = self.datatype_combo.currentText()
211
212        show_only_no_flag = self.no_flag_checkbox.isChecked()
213        show_only_with_flag = self.with_flag_checkbox.isChecked()
214
215        visible_count = 0
216        for i in range(self.file_list.count()):
217            item = self.file_list.item(i)
218            relative_path = item.text()
219            should_show = True
220
221            if search_text and search_text not in relative_path.lower():
222                should_show = False
223            if selected_subject != QCoreApplication.translate("Components", "All subjects") and selected_subject not in relative_path:
224                should_show = False
225            if selected_session != QCoreApplication.translate("Components", "All sessions") and selected_session not in relative_path:
226                should_show = False
227            if selected_modality != QCoreApplication.translate("Components", "All modalities") and selected_modality not in relative_path:
228                should_show = False
229            if selected_datatype != QCoreApplication.translate("Components", "All types") and f"/{selected_datatype}/" not in relative_path and f"\\{selected_datatype}\\" not in relative_path:
230                should_show = False
231
232            has_flag = relative_path in self.files_with_flag
233            if show_only_no_flag and has_flag:
234                should_show = False
235            if show_only_with_flag and not has_flag:
236                should_show = False
237
238            item.setHidden(not should_show)
239            if should_show:
240                visible_count += 1
241
242        self._update_info_label(visible_count)
243
244    # === Signals ===
245    def _connect_signals(self):
246        self.search_bar.textChanged.connect(self._apply_filters)
247        self.subject_combo.currentTextChanged.connect(self._apply_filters)
248        self.session_combo.currentTextChanged.connect(self._apply_filters)
249        self.modality_combo.currentTextChanged.connect(self._apply_filters)
250        self.datatype_combo.currentTextChanged.connect(self._apply_filters)
251        self.no_flag_checkbox.toggled.connect(self._apply_filters)
252        self.with_flag_checkbox.toggled.connect(self._apply_filters)
253
254        def on_no_flag_toggled(checked):
255            if checked:
256                self.with_flag_checkbox.setChecked(False)
257        def on_with_flag_toggled(checked):
258            if checked:
259                self.no_flag_checkbox.setChecked(False)
260        self.no_flag_checkbox.toggled.connect(on_no_flag_toggled)
261        self.with_flag_checkbox.toggled.connect(on_with_flag_toggled)
262
263        self._reset_button.clicked.connect(self._reset_filters)
264        self._select_all_button.clicked.connect(self._select_all_visible)
265        self._deselect_all_button.clicked.connect(lambda: self.file_list.clearSelection())
266
267        self.buttons.accepted.connect(self._accept)
268        self.buttons.rejected.connect(self.reject)
269
270    def _reset_filters(self):
271        self.search_bar.clear()
272        self.subject_combo.setCurrentIndex(0)
273        self.session_combo.setCurrentIndex(0)
274        self.modality_combo.setCurrentIndex(0)
275        self.datatype_combo.setCurrentIndex(0)
276        self.no_flag_checkbox.setChecked(False)
277        self.with_flag_checkbox.setChecked(False)
278
279    def _select_all_visible(self):
280        for i in range(self.file_list.count()):
281            item = self.file_list.item(i)
282            if not item.isHidden():
283                item.setSelected(True)
284
285    def _accept(self):
286        selected_items = [item for item in self.file_list.selectedItems() if not item.isHidden()]
287        if not selected_items:
288            QMessageBox.warning(
289                self,
290                QCoreApplication.translate("Components", "No selection"),
291                QCoreApplication.translate("Components", "Please select at least one visible NIfTI file."))
292            log.info("No selection")
293            return
294
295        selected_files = []
296        warnings = []
297
298        for item in selected_items:
299            relative_path = item.text()
300            abs_path = self.relative_to_absolute[relative_path]
301            path_parts = abs_path.replace(self.workspace_path, '').strip(os.sep).split(os.sep)
302            subject = next((p for p in path_parts if p.startswith('sub-')), None)
303            session = next((p for p in path_parts if p.startswith('ses-')), None)
304            filename = os.path.basename(abs_path)
305            filename_noext = os.path.splitext(os.path.splitext(filename)[0])[0]
306            modality = filename_noext.split("_")[-1] or QCoreApplication.translate("Components", "Unknown")
307            datatype = next((p for p in path_parts if p in ["anat", "func", "dwi", "fmap", "perf"]), None)
308
309            has_flag = relative_path in self.files_with_flag
310            if has_flag:
311                warnings.append((abs_path, subject))
312
313            selected_files.append(abs_path)
314
315        # Show warnings if needed
316        if warnings:
317            unique_subjects = sorted({s for _, s in warnings if s})
318            msg = QMessageBox(self)
319            msg.setIcon(QMessageBox.Icon.Warning)
320            msg.setWindowTitle(QCoreApplication.translate("Components", "Existing {0} Detected").format(self.label.title()))
321            if len(unique_subjects) == 1:
322                subject_display = unique_subjects[0] or QCoreApplication.translate("Components", "this patient")
323                msg.setText(QCoreApplication.translate("Components", "A {0} already exists for {1}.").format(self.label, subject_display))
324            else:
325                msg.setText(QCoreApplication.translate("Components", "{0}s already exist for {1} patients: {2}").format(self.label.title(), len(unique_subjects), ', '.join(unique_subjects)))
326            msg.setInformativeText(QCoreApplication.translate("Components", "You can still proceed to create additional {0}s. Continue?").format(self.label))
327            msg.setStandardButtons(QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No)
328            if msg.exec() == QMessageBox.StandardButton.No:
329                return
330
331        self.selected_files = selected_files
332        self.accept()
333
334    def _apply_forced_filters(self):
335        """Applica i filtri forzati definiti in self.forced_filters."""
336        if "search" in self.forced_filters:
337            self.search_bar.setText(self.forced_filters["search"])
338            self.search_bar.setEnabled(False)
339
340        def safe_set_combo_value(combo, key):
341            if key in self.forced_filters:
342                value = self.forced_filters[key]
343                index = combo.findText(value)
344                if index >= 0:
345                    combo.setCurrentIndex(index)
346                else:
347                    combo.addItem(value)
348                    combo.setCurrentIndex(combo.count() - 1)
349                combo.setEnabled(False)
350
351        safe_set_combo_value(self.subject_combo, "subject")
352        safe_set_combo_value(self.session_combo, "session")
353        safe_set_combo_value(self.modality_combo, "modality")
354        safe_set_combo_value(self.datatype_combo, "datatype")
355
356        if "no_flag" in self.forced_filters:
357            self.no_flag_checkbox.setChecked(bool(self.forced_filters["no_flag"]))
358            self.no_flag_checkbox.setEnabled(False)
359
360        if "with_flag" in self.forced_filters:
361            self.with_flag_checkbox.setChecked(bool(self.forced_filters["with_flag"]))
362            self.with_flag_checkbox.setEnabled(False)
363
364        # Riapplica i filtri
365        self._apply_filters()