Coverage for main/threads/skull_strip_thread.py: 93%

1import os 

2import json 

3import shutil 

4import tempfile 

5 

6from PyQt6.QtCore import pyqtSignal, QThread, QProcess, QCoreApplication 

7from logger import get_logger 

8from utils import setup_fsl_env, get_bin_path 

9 

10log = get_logger() 

11 

12 

13class SkullStripThread(QThread): 

14 """ 

15 Background worker thread for performing skull-stripping operations on NIfTI files. 

16 

17 This thread executes either **FSL BET** or **HD-BET** commands in the background using 

18 a `QProcess`, preventing the GUI from freezing during long-running neuroimaging operations. 

19 

20 It supports batch processing of multiple files, progress reporting, cancellation, and 

21 automatic generation of BIDS-like JSON sidecar metadata. 

22 

23 --- 

24 **Signals** 

25 - `progress_updated (str)`: Emits a human-readable progress message. 

26 - `progress_value_updated (int)`: Emits a numeric progress percentage. 

27 - `file_started (str)`: Emitted when processing of a specific file begins. 

28 - `file_completed (str, bool, str)`: Emitted when a file finishes (filename, success flag, message). 

29 - `all_completed (int, list)`: Emitted after all files are processed (success count, list of failed files). 

30 

31 --- 

32 

33 **Parameters** 

34 - `files (list[str])`: List of NIfTI file paths to process. 

35 - `workspace_path (str)`: Root directory of the workspace, containing derivatives. 

36 - `parameters (dict)`: Dictionary of BET/HD-BET options (e.g., `f_val`, `opt_m`, `opt_t`). 

37 - `has_cuda (bool)`: Whether CUDA GPU acceleration is available (for HD-BET). 

38 - `has_bet (bool)`: Whether to use FSL BET instead of HD-BET. 

39 """ 

40 

41 # Signal emitted when a progress message is updated. 

42 progress_updated = pyqtSignal(str) 

43 """**Signal(str):** Emitted when the progress message text changes.""" 

44 

45 # Signal emitted when the progress bar value changes. 

46 progress_value_updated = pyqtSignal(int) 

47 """**Signal(int):** Emitted when the progress value is updated.""" 

48 

49 # Signal emitted when processing of an individual file starts. 

50 file_started = pyqtSignal(str) 

51 """**Signal(str):** Emitted when processing of a file begins. The parameter is the file path or name.""" 

52 

53 # Signal emitted when processing of an individual file completes. 

54 file_completed = pyqtSignal(str, bool, str) 

55 """**Signal(str, bool, str):** Emitted when a file has finished processing. 

56 The parameters are: file path/name, success status (True if completed successfully), and an optional message or error description. 

57 """ 

58 

59 # Signal emitted when all files have completed processing. 

60 all_completed = pyqtSignal(int, list) 

61 """**Signal(int, list):** Emitted when all files are processed. 

62 The parameters are: the total number of processed files and a list of results or statuses. 

63 """ 

64 

65 def __init__(self, files, workspace_path, parameters, has_cuda, bet_tool): 

66 super().__init__() 

67 self.files = files 

68 self.workspace_path = workspace_path 

69 self.parameters = parameters 

70 self.has_cuda = has_cuda 

71 

72 self.bet_tool = bet_tool if bet_tool in ["fsl-bet","synthstrip","hd-bet"] else None 

73 self.is_cancelled = False 

74 self.success_count = 0 

75 self.failed_files = [] 

76 

77 def cancel(self): 

78 """ 

79 Cancels the ongoing operation. 

80 

81 If a subprocess (`QProcess`) is currently running, it is terminated and killed. 

82 This allows for safe interruption of long-running skull-stripping commands. 

83 """ 

84 self.is_cancelled = True 

85 

86 if hasattr(self, "process") and self.process is not None: 

87 try: 

88 self.process.terminate() 

89 self.process.kill() 

90 except Exception as e: 

91 log.error(f"Error while cancelling process: {e}") 

92 

93 def run(self): 

94 """ 

95 Executes the skull-stripping process for all input files. 

96 

97 For each file: 

98 1. Emits a `file_started` signal. 

99 2. Determines subject ID from the file path. 

100 3. Chooses between **BET** or **HD-BET** and builds the command. 

101 4. Runs the command asynchronously using `QProcess`. 

102 5. Writes the stripped output and corresponding JSON metadata. 

103 6. Emits `file_completed` for each processed file. 

104 

105 At the end, emits `all_completed` with a summary of successful and failed files. 

106 """ 

107 

108 if not self.files: 

109 log.warning("Nessun file da processare (lista vuota).") 

110 self.progress_value_updated.emit(100) 

111 self.progress_updated.emit( 

112 QCoreApplication.translate("Threads", "No files to process.") 

113 ) 

114 self.all_completed.emit(0, []) 

115 return 

116 

117 base_progress = 10 

118 progress_per_file = int(90 / len(self.files)) 

119 

120 for i, nifti_file in enumerate(self.files): 

121 if self.is_cancelled: 

122 break 

123 

124 filename = os.path.basename(nifti_file) 

125 self.file_started.emit(filename) 

126 

127 base_progress = int(base_progress + i * progress_per_file) 

128 self.progress_value_updated.emit(base_progress) 

129 self.progress_updated.emit( 

130 QCoreApplication.translate( 

131 "Threads", 

132 "Processing {0} ({1}/{2})" 

133 ).format(filename, i + 1, len(self.files)) 

134 ) 

135 

136 try: 

137 nifti_file = os.path.normpath(nifti_file) 

138 self.workspace_path = os.path.normpath(self.workspace_path) 

139 

140 # Extract subject ID from BIDS-like file path 

141 path_parts = nifti_file.replace(self.workspace_path, '').strip(os.sep).split(os.sep) 

142 subject_id = next((p for p in path_parts if p.startswith("sub-")), None) 

143 if not subject_id: 

144 self.file_completed.emit( 

145 filename, 

146 False, 

147 QCoreApplication.translate("Threads", "Cannot extract subject ID") 

148 ) 

149 self.failed_files.append(nifti_file) 

150 continue 

151 

152 # Define output directory for this subject 

153 output_dir = os.path.join( 

154 self.workspace_path, 

155 'derivatives', 

156 'skullstrips', 

157 subject_id, 

158 'anat' 

159 ) 

160 os.makedirs(output_dir, exist_ok=True) 

161 

162 base_name = filename.replace('.nii.gz', '').replace('.nii', '') 

163 

164 # Create an isolated temporary working directory 

165 temp_dir = tempfile.mkdtemp(prefix="skullstrip_") 

166 

167 # Build the appropriate command (BET or HD-BET) 

168 if self.bet_tool == "fsl-bet": 

169 # Configure FSL environment variables 

170 os.environ["FSLDIR"], os.environ["FSLOUTPUTTYPE"] = setup_fsl_env() 

171 f_val = self.parameters.get('f_val', 0.5) 

172 f_str = f"f{str(f_val).replace('.', '')}" 

173 

174 temp_output = os.path.join(temp_dir, f"{base_name}_{f_str}_brain.nii.gz") 

175 final_output = os.path.join(output_dir, f"{base_name}_{f_str}_brain.nii.gz") 

176 

177 cmd = ["bet", nifti_file, temp_output, "-f", str(f_val)] 

178 

179 # Optional BET parameters 

180 for opt in ['opt_m', 'opt_t', 'opt_s', 'opt_o']: 

181 if self.parameters.get(opt, False): 

182 cmd.append(f"-{opt[-1]}") 

183 

184 # Optional center coordinates 

185 for coord in ['c_x', 'c_y', 'c_z']: 

186 val = self.parameters.get(coord, 0) 

187 if val != 0: 

188 cmd += ["-c", str(val)] 

189 

190 # Exclude brain extraction mask if requested 

191 if not self.parameters.get('opt_brain_extracted', True): 

192 cmd.append("-n") 

193 

194 method = "FSL BET" 

195 

196 else: 

197 temp_output = os.path.join(temp_dir, f"{base_name}_{self.bet_tool}_brain.nii.gz") 

198 final_output = os.path.join(output_dir, f"{base_name}_{self.bet_tool}_brain.nii.gz") 

199 if self.bet_tool == "synthstrip": 

200 

201 cmd = [get_bin_path("mri_synthstrip"), "-i", nifti_file, "-o", temp_output] 

202 

203 method = "SynthStrip" 

204 

205 elif self.bet_tool == "hd-bet": 

206 

207 cmd = [get_bin_path("hd-bet"), "-i", nifti_file, "-o", temp_output] 

208 

209 # Disable CUDA if unavailable 

210 if not self.has_cuda: 

211 cmd += ["-device", "cpu", "--disable_tta"] 

212 

213 method = "HD-BET" 

214 

215 else: 

216 log.error("Bet tool not recognized.") 

217 return 

218 

219 # Run skull-stripping command 

220 self.process = QProcess() 

221 self.process.start(cmd[0], cmd[1:]) 

222 

223 # Wait for process completion, checking for cancellation 

224 while not self.process.waitForFinished(-1): 

225 if self.is_cancelled: 

226 self.process.kill() 

227 self.process.waitForFinished() 

228 break 

229 

230 # Handle cancellation mid-run 

231 if self.is_cancelled: 

232 shutil.rmtree(temp_dir, ignore_errors=True) 

233 break 

234 

235 # Retrieve exit information and logs 

236 ret_code = self.process.exitCode() 

237 stderr = bytes(self.process.readAllStandardError()).decode() 

238 stdout = bytes(self.process.readAllStandardOutput()).decode() 

239 

240 # Handle process errors 

241 if ret_code != 0: 

242 self.file_completed.emit( 

243 filename, 

244 False, 

245 stderr or QCoreApplication.translate("Threads", "Error executing command") 

246 ) 

247 self.failed_files.append(nifti_file) 

248 shutil.rmtree(temp_dir, ignore_errors=True) 

249 continue 

250 

251 # Move temporary output to final destination 

252 shutil.move(temp_output, final_output) 

253 shutil.rmtree(temp_dir, ignore_errors=True) 

254 

255 # Create BIDS-like JSON metadata 

256 json_file = final_output.replace(".nii.gz", ".json") 

257 metadata = { 

258 "SkullStripped": True, 

259 "Description": "Skull-stripped brain image", 

260 "Sources": [filename], 

261 "SkullStrippingMethod": method 

262 } 

263 

264 with open(json_file, 'w') as f: 

265 json.dump(metadata, f, indent=2) 

266 

267 # Success count and signal 

268 self.success_count += 1 

269 self.file_completed.emit(filename, True, "") 

270 

271 except Exception as e: 

272 # Catch all exceptions per file 

273 self.file_completed.emit(filename, False, str(e)) 

274 self.failed_files.append(nifti_file) 

275 try: 

276 shutil.rmtree(temp_dir, ignore_errors=True) 

277 except Exception: 

278 pass 

279 

280 # Emit final summary signal 

281 self.all_completed.emit(self.success_count, self.failed_files)