Coverage for main/components/nifti_file_dialog.py: 98%

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() 

366 

367 

368