fix to fold and memory leaks

This commit is contained in:
2026-01-29 22:30:28 -08:00
parent f82978e2e8
commit 92973da658
4 changed files with 177 additions and 168 deletions

310
main.py
View File

@@ -226,48 +226,6 @@ SECTIONS = [
from concurrent.futures import ProcessPoolExecutor
from PySide6.QtCore import QObject
from fold import run_fold_process
class FoldWorker(QObject):
progress_sig = Signal(int)
finished_sig = Signal(dict)
error_sig = Signal(str)
def __init__(self, haemo_obj, label, shared_dict):
super().__init__()
self.haemo_obj = haemo_obj
self.label = label
self.shared_dict = shared_dict
def run(self):
try:
with ProcessPoolExecutor(max_workers=1) as executor:
# Submit the function from the external file
future = executor.submit(run_fold_process, self.haemo_obj, self.label, self.shared_dict)
while not future.done():
current_progress = self.shared_dict.get(self.label, 0)
self.progress_sig.emit(current_progress)
time.sleep(0.1)
result = future.result()
# Check if our worker returned an error string instead of a dict
if isinstance(result, str) and result.startswith("ERROR:"):
raise Exception(result)
self.progress_sig.emit(100)
self.finished_sig.emit(result)
except Exception as e:
self.error_sig.emit(str(e))
class TerminalWindow(QWidget):
def __init__(self, parent=None):
super().__init__(parent, Qt.WindowType.Window)
@@ -1884,6 +1842,7 @@ class FullClickComboBox(QComboBox):
class ParticipantViewerWidget(QWidget):
def __init__(self, haemo_dict, fig_bytes_dict):
super().__init__()
self.setAttribute(Qt.WidgetAttribute.WA_DeleteOnClose)
self.setWindowTitle("FLARES Participant Viewer")
self.haemo_dict = haemo_dict
self.fig_bytes_dict = fig_bytes_dict
@@ -2044,10 +2003,11 @@ class ParticipantViewerWidget(QWidget):
def show_selected_images(self):
# Clear previous images
for i in reversed(range(self.grid_layout.count())):
widget = self.grid_layout.itemAt(i).widget()
while self.grid_layout.count():
item = self.grid_layout.takeAt(0)
widget = item.widget()
if widget:
widget.setParent(None)
widget.deleteLater()
selected_display_names = self._get_checked_items(self.participant_dropdown)
# Map from display names back to file paths
@@ -2140,7 +2100,9 @@ class ClickableLabel(QLabel):
self.setStyleSheet("border: 1px solid gray; margin: 2px;")
def mousePressEvent(self, event):
#TODO: This will use 3MB or RAM for every image that gets opened, and this RAM is not cleared when the expanded view is closed but only when the parent gets closed.
viewer = QWidget()
viewer.setAttribute(Qt.WidgetAttribute.WA_DeleteOnClose)
viewer.setWindowTitle("Expanded View")
layout = QVBoxLayout(viewer)
label = QLabel()
@@ -2498,6 +2460,23 @@ class MultiProgressDialog(QDialog):
def update_bar(self, label, value):
if label in self.bars:
self.bars[label].setValue(value)
def single_participant_worker(file_path, raw_data, result_queue, progress_queue):
""" Runs inside its own dedicated process """
p_name = os.path.basename(file_path)
try:
import flares
# Perform the heavy fold_channels logic
png_bytes = flares.fold_channels(raw_data)
# Hand back results and signal completion
result_queue.put({file_path: png_bytes})
progress_queue.put(p_name)
except Exception as e:
progress_queue.put(f"ERROR: {p_name} - {str(e)}")
class ParticipantFoldChannelsWidget(QWidget):
@@ -2506,7 +2485,6 @@ class ParticipantFoldChannelsWidget(QWidget):
self.setWindowTitle("FLARES Participant Fold Channels Viewer")
self.haemo_dict = haemo_dict
self.cha_dict = cha_dict
self.active_threads = []
# Create mappings: file_path -> participant label and dropdown display text
self.participant_map = {} # file_path -> "Participant 1"
self.participant_dropdown_items = [] # "Participant 1 (filename)"
@@ -2666,110 +2644,166 @@ class ParticipantFoldChannelsWidget(QWidget):
if not selected_display_names or 0 not in selected_indexes:
return
# 1. Setup the UI
self.progress_popup = MultiProgressDialog(self)
self.progress_popup.setWindowModality(Qt.WindowModality.NonModal) # Important!
self.active_threads = [] # Keep references alive
selected_files = [path for path, label in self.participant_map.items()
if f"{label} ({os.path.basename(path)})" in selected_display_names]
# 2. Create the Shared Memory Manager
# This allows the separate processes to "talk" to this GUI thread
self.process_manager = Manager()
self.shared_progress = self.process_manager.dict()
while self.grid_layout.count():
item = self.grid_layout.takeAt(0)
widget = item.widget()
if widget:
widget.deleteLater()
# 3. Launch Workers
for display_name in selected_display_names:
file_path = next((fp for fp, lbl in self.participant_map.items()
if f"{lbl} ({os.path.basename(fp)})" == display_name), None)
if not file_path: continue
# Use .copy() to ensure thread safety during pickling
haemo_obj = self.haemo_dict.get(file_path).copy()
label = self.participant_map[file_path]
# Initialize Shared Dict Entry
self.shared_progress[label] = 0
self.multi_progress = MultiProgressDialog(self)
for file_path in selected_files:
self.multi_progress.add_participant(os.path.basename(file_path), 1)
self.multi_progress.show()
# Add Bar to Popup
hbo_count = len(haemo_obj.copy().pick('hbo').ch_names) # Just for display logic if needed
self.progress_popup.add_participant(label, 100) # We use 0-100% standard
# Create Thread & Worker
thread = QThread()
# Pass the shared dict to the worker
worker = FoldWorker(haemo_obj, label, self.shared_progress)
worker.moveToThread(thread)
if current_process().name == 'MainProcess':
# Connect Signals
thread.started.connect(worker.run)
# 2. Setup Multiprocessing Manager
self.manager = Manager()
self.result_queue = self.manager.Queue()
self.progress_queue = self.manager.Queue()
self.active_processes = []
# 3. Start ALL processes at once
for file_path in selected_files:
p = Process(
target=single_participant_worker,
args=(file_path, self.haemo_dict[file_path], self.result_queue, self.progress_queue)
)
p.start()
self.active_processes.append(p)
# 4. Start the GUI listener
self.completed_count = 0
self.result_timer = QTimer()
self.result_timer.timeout.connect(self.check_parallel_results)
self.result_timer.start()
def check_parallel_results(self):
# Check for progress/completion signals
while not self.progress_queue.empty():
msg = self.progress_queue.get()
self.completed_count += 1
# Lambda capture to ensure correct label is used for each bar
worker.progress_sig.connect(lambda val, l=label: self.progress_popup.update_bar(l, val))
worker.finished_sig.connect(lambda imgs, l=label: self.on_fold_finished(imgs, l))
worker.error_sig.connect(self.on_fold_error)
# Cleanup Logic
worker.finished_sig.connect(thread.quit)
worker.finished_sig.connect(worker.deleteLater)
thread.finished.connect(thread.deleteLater)
# Store references to prevent garbage collection
self.active_threads.append({'thread': thread, 'worker': worker})
thread.start()
self.progress_popup.show()
def on_fold_finished(self, png_dict, label):
# 1. Close progress popup if all threads are done
# We filter the list to see if any threads are still running
still_running = any(t['thread'].isRunning() for t in self.active_threads)
if not still_running:
self.progress_popup.close()
# Optional: Shutdown manager when absolutely done to free resources
# self.process_manager.shutdown()
# 2. Display Images
if not hasattr(self, 'result_windows'):
self.result_windows = []
offset = len(self.result_windows) * 30 # Cascade offset
for key, png_data in png_dict.items():
popup = QDialog(self)
popup.setWindowTitle(f"{label} - fOLD {key.capitalize()}")
# ... (Your existing layout/image code) ...
# Resize and Position
if key == "main":
popup.resize(1664, 936)
popup.move(100 + offset, 100 + offset)
if msg.startswith("ERROR"):
print(f"Worker Error: {msg}")
else:
popup.resize(450, 800)
popup.move(1770 + offset, 100 + offset)
# msg is p_name here
self.multi_progress.update_bar(msg, 1)
# Pull images as they become available
while not self.result_queue.empty():
result_dict = self.result_queue.get()
self.add_images_to_grid(result_dict)
# Clean up when all processes are done
if self.completed_count >= len(self.active_processes):
self.result_timer.stop()
popup.show()
self.result_windows.append(popup)
# Close the custom multi-progress window
if hasattr(self, 'multi_progress'):
self.multi_progress.close()
# Clean up processes
for p in self.active_processes:
if p.is_alive():
p.join(timeout=1) # Give it a second to wrap up
p.close() # Explicitly close the process object
# Shut down the Manager process (the source of the 'rogue' process)
if hasattr(self, 'manager'):
self.manager.shutdown()
self.active_processes = []
print("Processing fully complete. All resources released.")
def add_images_to_grid(self, result_dict):
"""
result_dict format: { file_path: {"main": bytes, "legend": bytes} }
"""
for file_path, images in result_dict.items():
def on_fold_error(self, error_msg):
if hasattr(self, 'progress_popup'):
self.progress_popup.close()
msg_box = QMessageBox(self)
msg_box.setIcon(QMessageBox.Icon.Critical)
msg_box.setWindowTitle("Something went wrong!")
message = (
"Unable to locate the fOLD files!<br><br>"
f"Please download the 'Supplementary' folder from <a href='https://github.com/nirx/fOLD-public'>here</a>. "
"Once downloaded, place it in C:/Users/your username/mne_data/fOLD/fOLD-public-master/Supplementary."
)
msg_box.setTextFormat(Qt.TextFormat.RichText)
msg_box.setText(message)
msg_box.exec()
if self.grid_layout.count() == 0 and "legend" in images:
self._add_legend_to_grid(images["legend"])
# Create a container for this participant's results
container = QFrame()
container.setFrameShape(QFrame.StyledPanel)
vbox = QVBoxLayout(container)
participant_label = self.participant_map.get(file_path, os.path.basename(file_path))
title = QLabel(f"<b>{participant_label}</b>")
title.setAlignment(Qt.AlignCenter)
vbox.addWidget(title)
# We primarily want to show the 'main' plot in the grid
if "main" in images:
pixmap = self._bytes_to_pixmap(images["main"])
img_label = QLabel()
# Scale it to fit the thumbnail size defined in __init__
img_label.setPixmap(pixmap.scaled(
self.thumb_size,
Qt.KeepAspectRatio,
Qt.SmoothTransformation
))
img_label.setAlignment(Qt.AlignCenter)
# Optional: Click to open full size
img_label.mousePressEvent = lambda e, p=pixmap, t=participant_label: self._open_full_size(p, t)
vbox.addWidget(img_label)
# Determine grid position (row-major order)
count = self.grid_layout.count()
row = count // 3 # 3 columns wide
col = count % 3
self.grid_layout.addWidget(container, row, col)
def _bytes_to_pixmap(self, png_bytes):
"""Converts raw bytes from the multiprocess queue to a QPixmap."""
image = QImage.fromData(png_bytes)
return QPixmap.fromImage(image)
def _open_full_size(self, pixmap, title):
"""Simple popup to view the image at a readable scale."""
view = QDialog(self)
view.setWindowTitle(f"Full View - {title}")
layout = QVBoxLayout(view)
label = QLabel()
label.setPixmap(pixmap)
layout.addWidget(label)
view.show()
def _add_legend_to_grid(self, legend_bytes):
"""Helper to put the legend in the first slot."""
container = QFrame()
container.setStyleSheet("background-color: #f9f9f9; border: 1px solid #ccc;")
vbox = QVBoxLayout(container)
title = QLabel("<b>Brodmann Area Legend</b>")
title.setAlignment(Qt.AlignCenter)
vbox.addWidget(title)
pixmap = self._bytes_to_pixmap(legend_bytes)
legend_label = QLabel()
# Legends are usually tall, so we scale it differently or keep it smaller
legend_label.setPixmap(pixmap.scaled(
self.thumb_size,
Qt.KeepAspectRatio,
Qt.SmoothTransformation
))
legend_label.setAlignment(Qt.AlignCenter)
legend_label.mousePressEvent = lambda e, p=pixmap: self._open_full_size(p, "Brodmann Legend")
vbox.addWidget(legend_label)
self.grid_layout.addWidget(container, 0, 0)
class ExportDataAsCSVViewerWidget(QWidget):
def __init__(self, haemo_dict, cha_dict, df_ind, design_matrix, group, contrast_results_dict):