cleanup what I just did
This commit is contained in:
@@ -0,0 +1,287 @@
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import pandas as pd
|
||||
import numpy as np
|
||||
|
||||
from PySide6.QtWidgets import QWidget, QVBoxLayout, QLabel, QHBoxLayout, QMessageBox, QLineEdit, QPushButton, QFileDialog
|
||||
from PySide6.QtCore import Qt
|
||||
|
||||
from mne.io import read_raw_snirf
|
||||
from mne_nirs.io import write_raw_snirf
|
||||
|
||||
from src.shared.shareddata import APP_NAME
|
||||
|
||||
|
||||
class UpdateOptodesWindow(QWidget):
|
||||
|
||||
def __init__(self, parent=None):
|
||||
super().__init__(parent, Qt.WindowType.Window)
|
||||
self.setWindowTitle("Update optode positions")
|
||||
self.resize(760, 200)
|
||||
|
||||
self.label_file_a = QLabel("SNIRF file:")
|
||||
self.line_edit_file_a = QLineEdit()
|
||||
self.line_edit_file_a.setReadOnly(True)
|
||||
self.btn_browse_a = QPushButton("Browse .snirf")
|
||||
self.btn_browse_a.clicked.connect(self.browse_file_a)
|
||||
|
||||
self.label_file_b = QLabel("Text file:")
|
||||
self.line_edit_file_b = QLineEdit()
|
||||
self.line_edit_file_b.setReadOnly(True)
|
||||
self.btn_browse_b = QPushButton("Browse .txt/.xlsx")
|
||||
self.btn_browse_b.clicked.connect(self.browse_file_b)
|
||||
|
||||
self.label_suffix = QLabel("Suffix to append to filename:")
|
||||
self.line_edit_suffix = QLineEdit()
|
||||
self.line_edit_suffix.setText("flare")
|
||||
|
||||
self.btn_clear = QPushButton("Clear")
|
||||
self.btn_go = QPushButton("Go")
|
||||
self.btn_clear.clicked.connect(self.clear_files)
|
||||
self.btn_go.clicked.connect(self.go_action)
|
||||
|
||||
# ---
|
||||
layout = QVBoxLayout()
|
||||
self.description = QLabel()
|
||||
self.description.setTextFormat(Qt.TextFormat.RichText)
|
||||
self.description.setTextInteractionFlags(Qt.TextInteractionFlag.TextBrowserInteraction)
|
||||
self.description.setOpenExternalLinks(False) # Handle the click internally
|
||||
|
||||
self.description.setText("Some software when creating snirf files will insert a template of optode positions as the correct position of the optodes for the participant.<br>"
|
||||
"This is rarely correct as each head differs slightly in shape or size, and a lot of calculations require the optodes to be in the correct location.<br>"
|
||||
"Using a .txt or .xlsx file, we can update the positions in the snirf file to match those of a digitization system such as one from Polhemus or elsewhere.<br>"
|
||||
"The .txt file should have the fiducials, detectors, and sources clearly labeled, followed by the x, y, and z coordinates seperated by a space.<br>"
|
||||
"An example format of what a digitization text file should look like can be found <a href='custom_link'>by clicking here</a>. Currently only .xlsx files directly exported from a<br>"
|
||||
"Polhemus system are supported.")
|
||||
|
||||
self.description.linkActivated.connect(self.handle_link_click)
|
||||
layout.addWidget(self.description)
|
||||
|
||||
help_text_a = "Select the SNIRF (.snirf) file to update with new optode positions."
|
||||
|
||||
file_a_layout = QHBoxLayout()
|
||||
|
||||
# Help button on the left
|
||||
help_btn_a = QPushButton("?")
|
||||
help_btn_a.setFixedWidth(25)
|
||||
help_btn_a.setToolTip(help_text_a)
|
||||
help_btn_a.clicked.connect(lambda _, text=help_text_a: self.show_help_popup(text))
|
||||
file_a_layout.addWidget(help_btn_a)
|
||||
|
||||
# Container for label + line_edit + browse button with tooltip
|
||||
file_a_container = QWidget()
|
||||
file_a_container_layout = QHBoxLayout()
|
||||
file_a_container_layout.setContentsMargins(0, 0, 0, 0)
|
||||
file_a_container_layout.addWidget(self.label_file_a)
|
||||
file_a_container_layout.addWidget(self.line_edit_file_a)
|
||||
file_a_container_layout.addWidget(self.btn_browse_a)
|
||||
file_a_container.setLayout(file_a_container_layout)
|
||||
file_a_container.setToolTip(help_text_a)
|
||||
|
||||
file_a_layout.addWidget(file_a_container)
|
||||
layout.addLayout(file_a_layout)
|
||||
|
||||
help_text_b = "Provide a .txt file with labeled optodes (e.g., nz, rpa, lpa, d1, s1) and their x, y, z coordinates, or a .xlsx file from a Polhemius system."
|
||||
|
||||
file_b_layout = QHBoxLayout()
|
||||
|
||||
help_btn_b = QPushButton("?")
|
||||
help_btn_b.setFixedWidth(25)
|
||||
help_btn_b.setToolTip(help_text_b)
|
||||
help_btn_b.clicked.connect(lambda _, text=help_text_b: self.show_help_popup(text))
|
||||
file_b_layout.addWidget(help_btn_b)
|
||||
|
||||
file_b_container = QWidget()
|
||||
file_b_container_layout = QHBoxLayout()
|
||||
file_b_container_layout.setContentsMargins(0, 0, 0, 0)
|
||||
file_b_container_layout.addWidget(self.label_file_b)
|
||||
file_b_container_layout.addWidget(self.line_edit_file_b)
|
||||
file_b_container_layout.addWidget(self.btn_browse_b)
|
||||
file_b_container.setLayout(file_b_container_layout)
|
||||
file_b_container.setToolTip(help_text_b)
|
||||
|
||||
file_b_layout.addWidget(file_b_container)
|
||||
layout.addLayout(file_b_layout)
|
||||
|
||||
|
||||
help_text_suffix = "This text will be appended to the original filename when saving. Default is 'flare'."
|
||||
|
||||
suffix_layout = QHBoxLayout()
|
||||
|
||||
help_btn_suffix = QPushButton("?")
|
||||
help_btn_suffix.setFixedWidth(25)
|
||||
help_btn_suffix.setToolTip(help_text_suffix)
|
||||
help_btn_suffix.clicked.connect(lambda _, text=help_text_suffix: self.show_help_popup(text))
|
||||
suffix_layout.addWidget(help_btn_suffix)
|
||||
|
||||
suffix_container = QWidget()
|
||||
suffix_container_layout = QHBoxLayout()
|
||||
suffix_container_layout.setContentsMargins(0, 0, 0, 0)
|
||||
suffix_container_layout.addWidget(self.label_suffix)
|
||||
suffix_container_layout.addWidget(self.line_edit_suffix)
|
||||
suffix_container.setLayout(suffix_container_layout)
|
||||
suffix_container.setToolTip(help_text_suffix)
|
||||
|
||||
suffix_layout.addWidget(suffix_container)
|
||||
layout.addLayout(suffix_layout)
|
||||
|
||||
buttons_layout = QHBoxLayout()
|
||||
buttons_layout.addStretch()
|
||||
buttons_layout.addWidget(self.btn_clear)
|
||||
buttons_layout.addWidget(self.btn_go)
|
||||
layout.addLayout(buttons_layout)
|
||||
|
||||
self.setLayout(layout)
|
||||
|
||||
def show_help_popup(self, text):
|
||||
msg = QMessageBox(self)
|
||||
msg.setWindowTitle(f"Parameter Info - {APP_NAME.upper}")
|
||||
msg.setText(text)
|
||||
msg.exec()
|
||||
|
||||
def handle_link_click(self, link):
|
||||
if link == "custom_link":
|
||||
msg = QMessageBox(self)
|
||||
msg.setWindowTitle("Example Digitization File")
|
||||
|
||||
text = "nz: -1.91 85.175 -31.1525\n" \
|
||||
"rpa: 80.3825 -17.1925 -57.2775\n" \
|
||||
"lpa: -81.815 -17.1925 -57.965\n" \
|
||||
"d1: 0.01 -97.5175 62.5875\n" \
|
||||
"d2: 25.125 -103.415 45.045\n" \
|
||||
"d3: 49.095 -97.9025 30.2075\n" \
|
||||
"s1: 0.01 -112.43 32.595\n" \
|
||||
"s2: 30.325 -84.3125 71.8975\n" \
|
||||
"s3: 0.01 -70.6875 89.0925\n"
|
||||
msg.setText(text)
|
||||
msg.exec()
|
||||
|
||||
def browse_file_a(self):
|
||||
file_path, _ = QFileDialog.getOpenFileName(self, "Select SNIRF File", "", "SNIRF Files (*.snirf)")
|
||||
if file_path:
|
||||
self.line_edit_file_a.setText(file_path)
|
||||
|
||||
def browse_file_b(self):
|
||||
file_path, _ = QFileDialog.getOpenFileName(self, "Select File", "", "Supported Files (*.txt *.xlsx)")
|
||||
if file_path:
|
||||
self.line_edit_file_b.setText(file_path)
|
||||
|
||||
def clear_files(self):
|
||||
self.line_edit_file_a.clear()
|
||||
self.line_edit_file_b.clear()
|
||||
|
||||
def go_action(self):
|
||||
file_a = self.line_edit_file_a.text()
|
||||
file_b = self.line_edit_file_b.text()
|
||||
suffix = self.line_edit_suffix.text().strip() or "flare"
|
||||
|
||||
if not file_a:
|
||||
QMessageBox.critical(self, "Missing File", "Please select a SNIRF file.")
|
||||
return
|
||||
if not file_b:
|
||||
QMessageBox.critical(self, "Missing File", "Please select a TXT file.")
|
||||
return
|
||||
|
||||
# Get original filename without extension
|
||||
base_name = os.path.splitext(os.path.basename(file_a))[0]
|
||||
suggested_name = f"{base_name}_{suffix}.snirf"
|
||||
|
||||
# Open save dialog with default name
|
||||
save_path, _ = QFileDialog.getSaveFileName(
|
||||
self,
|
||||
"Save SNIRF File As",
|
||||
suggested_name,
|
||||
"SNIRF Files (*.snirf)"
|
||||
)
|
||||
|
||||
if not save_path:
|
||||
print("Save cancelled.")
|
||||
return
|
||||
|
||||
# Ensure .snirf extension
|
||||
if not save_path.lower().endswith(".snirf"):
|
||||
save_path += ".snirf"
|
||||
|
||||
try:
|
||||
self.update_optode_positions(file_a=file_a, file_b=file_b, save_path=save_path)
|
||||
except Exception as e:
|
||||
QMessageBox.critical(self, "Error", f"Failed to write file:\n{e}")
|
||||
return
|
||||
|
||||
QMessageBox.information(self, "File Saved", f"File was saved to:\n{save_path}")
|
||||
|
||||
def update_optode_positions(self, file_a, file_b, save_path):
|
||||
|
||||
fiducials = {}
|
||||
ch_positions = {}
|
||||
|
||||
extension = Path(file_b).suffix
|
||||
|
||||
# Read the lines from the optode file
|
||||
if extension == '.txt':
|
||||
with open(file_b, 'r') as f:
|
||||
for line in f:
|
||||
if line.strip():
|
||||
# Split by the semicolon and convert to meters
|
||||
ch_name, coords_str = line.split(":")
|
||||
coords = np.array(list(map(float, coords_str.strip().split()))) * 0.001
|
||||
|
||||
# The key we have is a fiducial
|
||||
if ch_name.lower() in ['lpa', 'nz', 'rpa']:
|
||||
fiducials[ch_name.lower()] = coords
|
||||
|
||||
# The key we have is a source or detector
|
||||
else:
|
||||
ch_positions[ch_name.upper()] = coords
|
||||
|
||||
elif extension == '.xlsx':
|
||||
|
||||
# TODO: Bad! Why assume sheet1 has the contents?
|
||||
df = pd.read_excel(file_b, sheet_name='Sheet1')
|
||||
|
||||
def _get_block_data(df, block_id, row_mapping, scale=0.001):
|
||||
"""Isolates a block, cleans numeric data, and returns a scaled dictionary."""
|
||||
# 1. Isolate and clean
|
||||
block = df[df['block_id'] == block_id].iloc[:, [1, 2, 3]].copy()
|
||||
block = block.apply(pd.to_numeric, errors='coerce')
|
||||
|
||||
# 2. Extract into dictionary based on mapping
|
||||
result = {}
|
||||
|
||||
# If row_mapping is a dict (like {0: 'nz'}), use it directly
|
||||
if isinstance(row_mapping, dict):
|
||||
for row_idx, key in row_mapping.items():
|
||||
if row_idx < len(block):
|
||||
result[key] = block.iloc[row_idx].to_numpy(dtype=float) * scale
|
||||
|
||||
# If row_mapping is a string prefix (like 'D' or 'S'), auto-generate keys
|
||||
elif isinstance(row_mapping, str):
|
||||
for i in range(len(block)):
|
||||
result[f"{row_mapping}{i+1}"] = block.iloc[i].to_numpy(dtype=float) * scale
|
||||
|
||||
return result
|
||||
|
||||
# Identify blocks
|
||||
is_empty = df.isnull().all(axis=1)
|
||||
df['block_id'] = is_empty.cumsum()
|
||||
clean_df = df[~is_empty].copy()
|
||||
|
||||
# Process Block 2: Landmarks
|
||||
fiducials = _get_block_data(clean_df, 2, {0: 'nz', 2: 'rpa', 3: 'lpa'})
|
||||
|
||||
# Process Block 3: D-Points
|
||||
d_points = _get_block_data(clean_df, 3, 'D')
|
||||
|
||||
# Process Block 4: S-Points
|
||||
s_points = _get_block_data(clean_df, 4, 'S')
|
||||
|
||||
ch_positions = {**d_points, **s_points}
|
||||
|
||||
# Create montage with updated coords in head space
|
||||
initial_montage = make_dig_montage(ch_pos=ch_positions, nasion=fiducials.get('nz'), lpa=fiducials.get('lpa'), rpa=fiducials.get('rpa'), coord_frame='head') # type: ignore
|
||||
|
||||
# Read the SNIRF file, set the montage, and write it back
|
||||
raw = read_raw_snirf(file_a, preload=True)
|
||||
raw.set_montage(initial_montage)
|
||||
write_raw_snirf(raw, save_path)
|
||||
Reference in New Issue
Block a user