Work on training and dataloading tasks

.gitignore

@@ -1,3 +1,4 @@
 .venv/
 __pycache__/
-.DS_Store
+.DS_Store
+outputs/

main.py

@@ -1,5 +1,6 @@
 import os
 import threading
+import gc
 from typing import Any, Dict
 
 import toml
@@ -15,20 +16,15 @@ from textual.widgets import (
     ProgressBar,
     RichLog,
     Rule,
-    Select,
-    TabbedContent,
-    TabPane,
 )
 
+from tasks import PIPELINE_TASKS, SCENARIOS
+from util.config_manager import handle_directory_config
 from util.progress import ProgressTracker
 from util.screens import OverwriteConfirmScreen
-from util.tasks import PIPELINE_TASKS, SCENARIOS
 from util.ui_logger import PipelineLogger
 
 
-# ==========================================
-# TUI Application
-# ==========================================
 class CNNHarnessApp(App):
     CSS_PATH = "style.tcss"
 
@@ -36,61 +32,43 @@ class CNNHarnessApp(App):
         super().__init__()
         self.stop_event = threading.Event()
         self.pause_event = threading.Event()
-
-        # Initialize our consolidated, thread-aware logger
+        self.current_config: Dict[str, Any] = {}
         self.ui_logger = PipelineLogger(self, widget_id="#live_log")
 
+    # ==========================================
+    # UI Layout
+    # ==========================================
     def compose(self) -> ComposeResult:
         yield Header()
 
         with Vertical(id="main_container"):
-            # --- TOP SECTION: Scrollable Configuration Tabs ---
-            with TabbedContent(initial="tab-general"):
-                with TabPane("General Config", id="tab-general"):
-                    with Horizontal(classes="row"):
-                        with Vertical(classes="input_group"):
-                            yield Label("Load Configuration:")
-                            yield Input(placeholder="./config.toml", id="config_path")
-                        yield Button("Load Config", variant="primary", id="load_btn")
-
-                    with Horizontal(classes="row"):
-                        with Vertical(classes="input_group"):
-                            yield Label("Save Directory:", classes="input_label")
-                            yield Input(
-                                placeholder="./models/experiment_1", id="save_dir"
-                            )
-                        with Vertical(classes="input_group"):
-                            yield Label("Batch Size:", classes="input_label")
-                            yield Input(placeholder="32", id="batch_size")
+            # --- TOP SECTION: Two-Column Configuration ---
+            with Horizontal(id="config_section"):
+                # Left Column: Path entry and buttons
+                with Vertical(id="config_left_col"):
+                    yield Label("Working Directory:")
+                    yield Input(placeholder="./outputs/experiment_1", id="work_dir")
+                    yield Button("Load / Init Dir", variant="primary", id="load_btn")
+                    yield Button("Reload Config", variant="default", id="reload_btn")
+
+                # Right Column: Read-only parameters display
+                with Vertical(id="config_right_col"):
+                    yield Label("Current Parameters:")
+                    yield RichLog(id="config_display", highlight=True, markup=True)
+
+            yield Rule()
 
             # --- BOTTOM SECTION: Controls, Progress, and Logs ---
             with Horizontal(id="bottom_area"):
+                # --- Pipeline Controls ---
                 with Vertical(id="control_panel"):
-                    with Horizontal(id="status_row"):
-                        scenario_options = [
-                            (data["label"], sc_id) for sc_id, data in SCENARIOS.items()
-                        ]
-                        yield Select(
-                            scenario_options,
-                            prompt="Select Pipeline Scenario...",
-                            id="scenario_select",
-                        )
-                        yield Label("STATUS: READY", id="pipeline_status_label")
-
-                    yield Rule()
-
-                    yield Button(
-                        "Start Pipeline & Save Config",
-                        variant="success",
-                        id="btn_start",
-                    )
+                    yield Button("Start Pipeline", variant="success", id="btn_start")
 
                     with Horizontal(id="active_controls"):
                         yield Button("Pause", variant="warning", id="btn_pause")
                         yield Button("Stop", variant="error", id="btn_stop")
 
-                    yield Rule()
-
+                    # --- Progress Trackers ---
                     with Vertical(id="progress_area"):
                         for i in range(5):
                             with Horizontal(
@@ -108,8 +86,12 @@ class CNNHarnessApp(App):
                                     id=f"progress_stats_{i}", classes="progress_stats"
                                 )
 
+                # --- Live Log Output ---
                 with Vertical(id="log_panel"):
-                    yield Label("Live Logs:")
+                    with Horizontal(id="log_status_row"):
+                        yield Label("STATUS: READY", id="pipeline_status_label")
+                    with Horizontal(id="log_header_row"):
+                        yield Label("Live Logs:", classes="section_label")
                     yield RichLog(id="live_log", highlight=True, markup=True)
 
         yield Footer()
@@ -118,98 +100,83 @@ class CNNHarnessApp(App):
         self.query_one("#active_controls").display = False
         self.ui_logger.info("Application initialized and ready.")
 
+    # ==========================================
+    # Helper Methods
+    # ==========================================
     def _set_pipeline_status(self, text: str) -> None:
         try:
             self.query_one("#pipeline_status_label", Label).update(text)
         except Exception:
             pass
 
-    # ==========================================
-    # Configuration Management Methods
-    # ==========================================
-
-    def _load_config_file(self, filepath: str) -> Dict[str, Any]:
-        with open(filepath, "r") as f:
-            return toml.load(f)
-
-    def _save_config_file(self, filepath: str, config: Dict[str, Any]) -> None:
-        os.makedirs(os.path.dirname(filepath), exist_ok=True)
-        with open(filepath, "w") as f:
-            toml.dump(config, f)
+    def _update_config_display(self) -> None:
+        config_log = self.query_one("#config_display", RichLog)
+        config_log.clear()
 
-    def _populate_ui_from_config(self, config: Dict[str, Any]) -> None:
-        self.query_one("#save_dir", Input).value = str(config.get("save_dir", ""))
-        self.query_one("#batch_size", Input).value = str(config.get("batch_size", ""))
+        if not self.current_config:
+            config_log.write("[italic]No configuration loaded.[/italic]")
+            return
 
-        loaded_scenario = config.get("scenario", "")
-        if loaded_scenario in SCENARIOS:
-            self.query_one("#scenario_select", Select).value = loaded_scenario
-
-    def _gather_config_from_ui(self) -> Dict[str, Any]:
-        scenario_select = self.query_one("#scenario_select", Select)
-
-        scenario_id = (
-            scenario_select.value if type(scenario_select.value) is str else None
-        )
-
-        config = {
-            "save_dir": self.query_one("#save_dir", Input).value or "./outputs/default",
-            "batch_size": self.query_one("#batch_size", Input).value or "32",
-            "scenario": scenario_id,
-        }
+        toml_string = toml.dumps(self.current_config)
+        config_log.write(toml_string)
 
-        if scenario_id and scenario_id in SCENARIOS:
-            config["steps"] = SCENARIOS[scenario_id]["tasks"]
-        else:
-            config["steps"] = []
-
-        return config
+    def _toggle_config_inputs(self, disabled: bool) -> None:
+        self.query_one("#work_dir", Input).disabled = disabled
+        self.query_one("#load_btn", Button).disabled = disabled
+        self.query_one("#reload_btn", Button).disabled = disabled
 
     # ==========================================
     # Event Handlers
     # ==========================================
-
     def on_button_pressed(self, event: Button.Pressed) -> None:
-        if event.button.id == "load_btn":
-            config_path = self.query_one("#config_path", Input).value
-            if not config_path:
-                self.ui_logger.error("Please specify a config path to load.")
-                return
 
-            try:
-                loaded_config = self._load_config_file(config_path)
-                self._populate_ui_from_config(loaded_config)
-                self.ui_logger.info(f"Successfully loaded config from {config_path}")
-            except Exception as e:
-                self.ui_logger.error(f"Failed to load config: {str(e)}")
+        if event.button.id in ("load_btn", "reload_btn"):
+            work_dir = self.query_one("#work_dir", Input).value
+            success, message, config_data = handle_directory_config(work_dir)
+
+            if success:
+                self.current_config = config_data
+                self.current_config["work_dir"] = work_dir
+                self._update_config_display()
+                self.ui_logger.info(message)
+            else:
+                self.ui_logger.error(message)
 
         elif event.button.id == "btn_start":
-            config_dict = self._gather_config_from_ui()
+            if not self.current_config:
+                self.ui_logger.error(
+                    "No configuration loaded! Please load a directory first."
+                )
+                return
 
-            if not config_dict.get("scenario"):
+            scenario_id = self.current_config.get("scenario")
+            if not scenario_id or scenario_id not in SCENARIOS:
                 self.ui_logger.error(
-                    "No scenario selected! Please select a scenario from the dropdown."
+                    f"Invalid or missing scenario '{scenario_id}' in config.toml!"
                 )
                 return
 
-            save_dir = config_dict["save_dir"]
+            work_dir = self.current_config["work_dir"]
 
-            if os.path.exists(save_dir) and os.listdir(save_dir):
+            # Check if the working directory exists and is non-empty (except for the config.toml file)
+            if os.path.exists(work_dir) and any(
+                f for f in os.listdir(work_dir) if f != "config.toml"
+            ):
 
                 def check_overwrite_callback(proceed: bool) -> None:
                     if proceed:
-                        self._execute_pipeline(config_dict)
+                        self._execute_pipeline()
                     else:
                         self.ui_logger.error(
                             "Pipeline start cancelled by user (folder non-empty)."
                         )
 
                 self.app.push_screen(
-                    OverwriteConfirmScreen(save_dir),
+                    OverwriteConfirmScreen(work_dir),
                     check_overwrite_callback,  # pyright: ignore
                 )
             else:
-                self._execute_pipeline(config_dict)
+                self._execute_pipeline()
 
         elif event.button.id == "btn_pause":
             if self.pause_event.is_set():
@@ -236,11 +203,8 @@ class CNNHarnessApp(App):
     # ==========================================
     # Pipeline Execution
     # ==========================================
-
-    def _execute_pipeline(self, config_dict: Dict[str, Any]) -> None:
-        self.query_one(TabbedContent).disabled = True
-        self.query_one("#scenario_select").disabled = True
-
+    def _execute_pipeline(self) -> None:
+        self._toggle_config_inputs(disabled=True)
         self.query_one("#btn_start").display = False
         self.query_one("#active_controls").display = True
 
@@ -248,31 +212,28 @@ class CNNHarnessApp(App):
 
         self.stop_event.clear()
         self.pause_event.clear()
-        self.query_one("#btn_pause", Button).label = "Pause"
-        self.query_one("#btn_pause", Button).variant = "warning"
 
-        self.run_background_pipeline(config_dict)
+        btn_pause = self.query_one("#btn_pause", Button)
+        btn_pause.label = "Pause"
+        btn_pause.variant = "warning"
+
+        self.run_background_pipeline(self.current_config.copy())
 
     @work(exclusive=True, thread=True)
     def run_background_pipeline(self, config: Dict[str, Any]) -> None:
-        save_dir = config["save_dir"]
-        steps = config["steps"]
-        scenario_label = SCENARIOS[config["scenario"]]["label"]
+        scenario_id = config["scenario"]
+        steps = SCENARIOS[scenario_id]["tasks"]
+        scenario_label = SCENARIOS[scenario_id]["label"]
+
         root_tracker = ProgressTracker(self, level=0)
+        pipeline_state: Dict[str, Any] = {
+            "events": {"stop": self.stop_event, "pause": self.pause_event}
+        }
 
         try:
-            config_out_path = os.path.join(save_dir, "config.toml")
-            self._save_config_file(config_out_path, config)
-
             self.ui_logger.info(f"Initiating Scenario: {scenario_label}")
-            self.ui_logger.info(f"Configuration saved to {config_out_path}")
-
             total_tasks = len(steps)
 
-            pipeline_state: Dict[str, Any] = {
-                "events": {"stop": self.stop_event, "pause": self.pause_event}
-            }
-
             root_tracker.reset()
             root_tracker.set_title("Pipeline Status")
             root_tracker.update(total=total_tasks, advance=0)
@@ -284,12 +245,11 @@ class CNNHarnessApp(App):
                 task_name = PIPELINE_TASKS[task_id]["task_name"]
                 task_func = PIPELINE_TASKS[task_id]["task_func"]
 
-                # Generate a task-specific logger from our PipelineLogger
                 task_logger = self.ui_logger.get_task_logger(task_name)
-
                 self.ui_logger.info(f"Starting Task {i + 1}/{total_tasks}: {task_name}")
 
                 task_tracker = root_tracker.get_sub_tracker(task_name)
+
                 task_func(task_tracker, task_logger, config, pipeline_state)
 
                 if not self.stop_event.is_set():
@@ -326,14 +286,15 @@ class CNNHarnessApp(App):
             )
 
         finally:
+            pipeline_state.clear()
+            gc.collect()
             self.call_from_thread(self._reset_ui)
 
     def _reset_ui(self) -> None:
         self.query_one("#btn_stop", Button).disabled = False
         self.query_one("#btn_start").display = True
         self.query_one("#active_controls").display = False
-        self.query_one(TabbedContent).disabled = False
-        self.query_one("#scenario_select").disabled = False
+        self._toggle_config_inputs(disabled=False)
 
 
 if __name__ == "__main__":

model/dataset.py

@@ -23,13 +23,14 @@ def _row_to_float_tensor(row: pd.DataFrame, *, image_id: int) -> torch.Tensor:
 
     return torch.tensor(numeric.to_numpy(dtype="float32"))  # pyright: ignore
 
+
 def xls_pre(df: pd.DataFrame) -> pd.DataFrame:
     """
     Preprocess the Excel DataFrame.
     This function can be customized to filter or modify the DataFrame as needed.
     """
 
-    data: pd.DataFrame = df[["Image Data ID", "Sex", "Age (current)"]] # pyright: ignore
+    data: pd.DataFrame = df[["Image Data ID", "Sex", "Age (current)"]]  # pyright: ignore
     data["Sex"] = data["Sex"].str.strip()  # type: ignore
     data = data.replace({"M": 0, "F": 1})  # type: ignore
     data.set_index("Image Data ID")  # type: ignore

model/training.py

@@ -1,4 +1,6 @@
 import pathlib as pl
+import time
+from threading import Event
 from typing import Callable, Tuple, cast
 
 import numpy as np
@@ -8,6 +10,8 @@ import torch.nn as nn
 from torch.utils.data import DataLoader
 
 from model.dataset import ADNIDataset
+from util.progress import ProgressTracker
+from util.ui_logger import PipelineLogger
 
 type TrainMetrics = Tuple[
     float, float, float, float
@@ -17,12 +21,47 @@ type TestMetrics = Tuple[float, float]  # (test_loss, test_acc)
 type KLLossFn = Callable[[nn.Module], torch.Tensor | None]
 
 
+def _batch_correct_and_total(
+    outputs: torch.Tensor,
+    targets: torch.Tensor,
+) -> Tuple[int, int]:
+    if outputs.ndim > 1 and outputs.size(-1) > 1:
+        predicted_classes = outputs.argmax(dim=1)
+        if targets.ndim > 1 and targets.size(-1) > 1:
+            target_classes = targets.argmax(dim=1)
+        else:
+            target_classes = targets.long().view(-1)
+
+        correct = int((predicted_classes == target_classes).sum().item())
+        total = int(target_classes.numel())
+        return correct, total
+
+    predicted = (outputs > 0.5).to(dtype=targets.dtype)
+    correct = int((predicted == targets).sum().item())
+    total = int(targets.numel())
+    return correct, total
+
+
+def _check_control_events(
+    stop_event: Event | None,
+    pause_event: Event | None,
+) -> None:
+    if stop_event is not None and stop_event.is_set():
+        raise InterruptedError("Pipeline execution stopped by user.")
+    while pause_event is not None and pause_event.is_set():
+        time.sleep(0.5)
+        if stop_event is not None and stop_event.is_set():
+            raise InterruptedError("Pipeline execution stopped by user while paused.")
+
+
 def test_model(
     model: nn.Module,
     test_loader: DataLoader[ADNIDataset],
     criterion: nn.Module,
-    progress: Callable[[int | None, int], None],
-    write_log: Callable[[str, bool], None],
+    progress: ProgressTracker,
+    log: PipelineLogger,
+    stop_event: Event | None = None,
+    pause_event: Event | None = None,
 ) -> TestMetrics:
     """
     Tests the model on the test dataset.
@@ -39,18 +78,25 @@ def test_model(
     test_loss = 0.0
     correct = 0
     total = 0
+    total_samples = 0
 
+    test_progress = progress.get_sub_tracker("Testing Batches")
+    test_progress.update(total=len(test_loader), advance=0)
     for _, (mri, xls, targets, _) in enumerate(test_loader):
+        _check_control_events(stop_event=stop_event, pause_event=pause_event)
         outputs = model((mri, xls))
         loss = criterion(outputs, targets)
-        test_loss += loss.item() * (mri.size(0) + xls.size(0))
+        batch_size = mri.size(0)
+        test_loss += loss.item() * batch_size
+        total_samples += batch_size
 
         # Calculate accuracy
-        predicted = (outputs > 0.5).float()
-        correct += (predicted == targets).sum().item()
-        total += targets.numel()
+        batch_correct, batch_total = _batch_correct_and_total(outputs, targets)
+        correct += batch_correct
+        total += batch_total
+        test_progress.update(total=len(test_loader), advance=1)
 
-    test_loss /= len(test_loader)
+    test_loss = test_loss / total_samples if total_samples > 0 else 0.0
     test_acc = correct / total if total > 0 else 0.0
     return test_loss, test_acc
 
@@ -61,6 +107,10 @@ def train_epoch(
     val_loader: DataLoader[ADNIDataset],
     optimizer: torch.optim.Optimizer,
     criterion: nn.Module,
+    log: PipelineLogger,
+    progress: ProgressTracker,
+    stop_event: Event | None = None,
+    pause_event: Event | None = None,
 ) -> Tuple[float, float, float, float]:
     """
     Trains the model for one epoch and evaluates it on the validation set.
@@ -77,36 +127,59 @@ def train_epoch(
     """
     model.train()
     train_loss = 0.0
-
-    # Training loop
+    train_correct = 0
+    train_total = 0
+    train_total_samples = 0
+
+    # Training loop - with progress bar for the training batches
+    train_progress = progress.get_sub_tracker("Training Batches")
+    train_progress.reset()
+    train_progress.set_title("Training Batches")
+    train_progress.update(total=len(train_loader), advance=0)
     for _, (mri, xls, targets, _) in enumerate(train_loader):
+        _check_control_events(stop_event=stop_event, pause_event=pause_event)
         optimizer.zero_grad()
         outputs = model((mri, xls))
         loss = criterion(outputs, targets)
         loss.backward()  # type: ignore[reportUnknownMemberType]
         optimizer.step()
-        train_loss += loss.item() * (mri.size(0) + xls.size(0))
-    train_loss /= len(train_loader)
+        batch_size = mri.size(0)
+        train_loss += loss.item() * batch_size
+        train_total_samples += batch_size
+        batch_correct, batch_total = _batch_correct_and_total(outputs, targets)
+        train_correct += batch_correct
+        train_total += batch_total
+        train_progress.update(total=len(train_loader), advance=1)
+    train_loss = train_loss / train_total_samples if train_total_samples > 0 else 0.0
 
     model.eval()
     val_loss = 0.0
     correct = 0
     total = 0
+    val_total_samples = 0
 
+    test_progress = progress.get_sub_tracker("Validation Batches")
+    test_progress.reset()
+    test_progress.set_title("Validation Batches")
+    test_progress.update(total=len(val_loader), advance=0)
     with torch.no_grad():
         for _, (mri, xls, targets, _) in enumerate(val_loader):
+            _check_control_events(stop_event=stop_event, pause_event=pause_event)
             outputs = model((mri, xls))
             loss = criterion(outputs, targets)
-            val_loss += loss.item() * (mri.size(0) + xls.size(0))
+            batch_size = mri.size(0)
+            val_loss += loss.item() * batch_size
+            val_total_samples += batch_size
 
             # Calculate accuracy
-            predicted = (outputs > 0.5).float()
-            correct += (predicted == targets).sum().item()
-            total += targets.numel()
+            batch_correct, batch_total = _batch_correct_and_total(outputs, targets)
+            correct += batch_correct
+            total += batch_total
+            test_progress.update(total=len(val_loader), advance=1)
 
-    val_loss /= len(val_loader)
+    val_loss = val_loss / val_total_samples if val_total_samples > 0 else 0.0
     val_acc = correct / total if total > 0 else 0.0
-    train_acc = correct / total if total > 0 else 0.0
+    train_acc = train_correct / train_total if train_total > 0 else 0.0
     return train_loss, val_loss, train_acc, val_acc
 
 
@@ -118,6 +191,10 @@ def train_model(
     criterion: nn.Module,
     num_epochs: int,
     output_path: pl.Path,
+    progress: ProgressTracker,
+    log: PipelineLogger,
+    stop_event: Event | None = None,
+    pause_event: Event | None = None,
 ) -> Tuple[nn.Module, pd.DataFrame]:
     """
     Trains the model using the provided training and validation data loaders.
@@ -139,13 +216,22 @@ def train_model(
 
     nhist = np.zeros((num_epochs, 4), dtype=np.float32)
 
+    # Get new progress bar for epochs
+
+    epoch_progress = progress.get_sub_tracker("Epochs")
+    epoch_progress.update(total=num_epochs, advance=0)
     for epoch in range(num_epochs):
+        _check_control_events(stop_event=stop_event, pause_event=pause_event)
         train_loss, val_loss, train_acc, val_acc = train_epoch(
             model,
             train_loader,
             val_loader,
             optimizer,
             criterion,
+            log=log,
+            progress=epoch_progress,
+            stop_event=stop_event,
+            pause_event=pause_event,
         )
 
         # Update the history
@@ -154,7 +240,7 @@ def train_model(
         nhist[epoch, 2] = train_acc
         nhist[epoch, 3] = val_acc
 
-        print(
+        log.info(
             f"Epoch [{epoch + 1}/{num_epochs}], "
             f"Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}, "
             f"Train Acc: {train_acc:.4f}, Val Acc: {val_acc:.4f}"
@@ -167,7 +253,9 @@ def train_model(
                     output_path / "intermediate_models" / f"model_epoch_{epoch + 1}.pt"
                 )
                 torch.save(model.state_dict(), model_save_path)
-                print(f"Model saved at epoch {epoch + 1}")
+                log.info(f"Model saved at epoch {epoch + 1}")
+
+        epoch_progress.update(total=num_epochs, advance=1)
 
     # return the trained model and the traning history
 
@@ -185,6 +273,10 @@ def test_model_bayesian(
     test_loader: DataLoader[ADNIDataset],
     criterion: nn.Module,
     get_kl_loss: KLLossFn,
+    progress: ProgressTracker,
+    log: PipelineLogger,
+    stop_event: Event | None = None,
+    pause_event: Event | None = None,
 ) -> TestMetrics:
     """
     Tests a Bayesian model on the test dataset with KL-augmented loss.
@@ -193,9 +285,13 @@ def test_model_bayesian(
     test_loss = 0.0
     correct = 0
     total = 0
+    total_samples = 0
 
+    test_progress = progress.get_sub_tracker("Testing Batches")
+    test_progress.update(total=len(test_loader), advance=0)
     with torch.no_grad():
         for _, (mri, xls, targets, _) in enumerate(test_loader):
+            _check_control_events(stop_event=stop_event, pause_event=pause_event)
             outputs = model((mri, xls))
             data_loss = cast(torch.Tensor, criterion(outputs, targets))
             batch_size = mri.size(0)
@@ -207,13 +303,15 @@ def test_model_bayesian(
             )
             loss: torch.Tensor = data_loss + kl_loss
 
-            test_loss += loss.item() * (mri.size(0) + xls.size(0))
+            test_loss += loss.item() * batch_size
+            total_samples += batch_size
 
-            predicted = (outputs > 0.5).float()
-            correct += (predicted == targets).sum().item()
-            total += targets.numel()
+            batch_correct, batch_total = _batch_correct_and_total(outputs, targets)
+            correct += batch_correct
+            total += batch_total
+            test_progress.update(total=len(test_loader), advance=1)
 
-    test_loss /= len(test_loader)
+    test_loss = test_loss / total_samples if total_samples > 0 else 0.0
     test_acc = correct / total if total > 0 else 0.0
     return test_loss, test_acc
 
@@ -225,14 +323,26 @@ def train_epoch_bayesian(
     optimizer: torch.optim.Optimizer,
     criterion: nn.Module,
     get_kl_loss: KLLossFn,
+    log: PipelineLogger,
+    progress: ProgressTracker,
+    stop_event: Event | None = None,
+    pause_event: Event | None = None,
 ) -> TrainMetrics:
     """
     Trains a Bayesian model for one epoch and evaluates on validation data.
     """
     model.train()
     train_loss = 0.0
-
+    train_correct = 0
+    train_total = 0
+    train_total_samples = 0
+
+    train_progress = progress.get_sub_tracker("Training Batches")
+    train_progress.reset()
+    train_progress.set_title("Training Batches")
+    train_progress.update(total=len(train_loader), advance=0)
     for _, (mri, xls, targets, _) in enumerate(train_loader):
+        _check_control_events(stop_event=stop_event, pause_event=pause_event)
         optimizer.zero_grad()
         outputs = model((mri, xls))
         data_loss = cast(torch.Tensor, criterion(outputs, targets))
@@ -246,16 +356,27 @@ def train_epoch_bayesian(
         loss: torch.Tensor = data_loss + kl_loss
         loss.backward()  # type: ignore[reportUnknownMemberType]
         optimizer.step()
-        train_loss += loss.item() * (mri.size(0) + xls.size(0))
-    train_loss /= len(train_loader)
+        train_loss += loss.item() * batch_size
+        train_total_samples += batch_size
+        batch_correct, batch_total = _batch_correct_and_total(outputs, targets)
+        train_correct += batch_correct
+        train_total += batch_total
+        train_progress.update(total=len(train_loader), advance=1)
+    train_loss = train_loss / train_total_samples if train_total_samples > 0 else 0.0
 
     model.eval()
     val_loss = 0.0
     correct = 0
     total = 0
+    val_total_samples = 0
 
+    test_progress = progress.get_sub_tracker("Validation Batches")
+    test_progress.reset()
+    test_progress.set_title("Validation Batches")
+    test_progress.update(total=len(val_loader), advance=0)
     with torch.no_grad():
         for _, (mri, xls, targets, _) in enumerate(val_loader):
+            _check_control_events(stop_event=stop_event, pause_event=pause_event)
             outputs = model((mri, xls))
             data_loss = cast(torch.Tensor, criterion(outputs, targets))
             batch_size = mri.size(0)
@@ -266,15 +387,17 @@ def train_epoch_bayesian(
                 else torch.tensor(0.0, device=outputs.device)
             )
             vloss: torch.Tensor = data_loss + kl_loss
-            val_loss += vloss.item() * (mri.size(0) + xls.size(0))
+            val_loss += vloss.item() * batch_size
+            val_total_samples += batch_size
 
-            predicted = (outputs > 0.5).float()
-            correct += (predicted == targets).sum().item()
-            total += targets.numel()
+            batch_correct, batch_total = _batch_correct_and_total(outputs, targets)
+            correct += batch_correct
+            total += batch_total
+            test_progress.update(total=len(val_loader), advance=1)
 
-    val_loss /= len(val_loader)
+    val_loss = val_loss / val_total_samples if val_total_samples > 0 else 0.0
     val_acc = correct / total if total > 0 else 0.0
-    train_acc = correct / total if total > 0 else 0.0
+    train_acc = train_correct / train_total if train_total > 0 else 0.0
 
     return train_loss, val_loss, train_acc, val_acc
 
@@ -288,13 +411,20 @@ def train_model_bayesian(
     num_epochs: int,
     output_path: pl.Path,
     get_kl_loss: KLLossFn,
+    progress: ProgressTracker,
+    log: PipelineLogger,
+    stop_event: Event | None = None,
+    pause_event: Event | None = None,
 ) -> Tuple[nn.Module, pd.DataFrame]:
     """
     Trains a Bayesian model with KL-augmented objective.
     """
     nhist = np.zeros((num_epochs, 4), dtype=np.float32)
 
+    epoch_progress = progress.get_sub_tracker("Epochs")
+    epoch_progress.update(total=num_epochs, advance=0)
     for epoch in range(num_epochs):
+        _check_control_events(stop_event=stop_event, pause_event=pause_event)
         train_loss, val_loss, train_acc, val_acc = train_epoch_bayesian(
             model,
             train_loader,
@@ -302,6 +432,10 @@ def train_model_bayesian(
             optimizer,
             criterion,
             get_kl_loss,
+            log=log,
+            progress=epoch_progress,
+            stop_event=stop_event,
+            pause_event=pause_event,
         )
 
         nhist[epoch, 0] = train_loss
@@ -309,7 +443,7 @@ def train_model_bayesian(
         nhist[epoch, 2] = train_acc
         nhist[epoch, 3] = val_acc
 
-        print(
+        log.info(
             f"Epoch [{epoch + 1}/{num_epochs}], "
             f"Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}, "
             f"Train Acc: {train_acc:.4f}, Val Acc: {val_acc:.4f}"
@@ -321,7 +455,9 @@ def train_model_bayesian(
                     output_path / "intermediate_models" / f"model_epoch_{epoch + 1}.pt"
                 )
                 torch.save(model.state_dict(), model_save_path)
-                print(f"Model saved at epoch {epoch + 1}")
+                log.info(f"Model saved at epoch {epoch + 1}")
+
+        epoch_progress.update(total=num_epochs, advance=1)
 
     history = pd.DataFrame(
         data=nhist.astype(np.float32),

outputs/default/config.toml

@@ -1,4 +0,0 @@
-save_dir = "./outputs/default"
-batch_size = "32"
-scenario = "scen_train_all"
-steps = [ "load_data", "train_regular", "train_bayesian", "evaluate_regular", "evaluate_bayesian",]

pyproject.toml

@@ -25,4 +25,4 @@ dependencies = [
 
 
 [tool.mypy]
-exclude = [".venv/**"]
+exclude = [".venv/**"]

style.tcss

@@ -1,25 +1,46 @@
 Screen { align: center middle; }
 #main_container { width: 95%; height: 95%; border: solid green; padding: 1 2; }
 
-/* Top Section: Tabs (Fixed percentage, scrollable) */
-TabbedContent {
-    height: 40%;
+/* =========================================
+   Top Section: Configuration
+   ========================================= */
+#config_section {
+    height: 35%;          /* Constrains height so it doesn't take over */
     margin-bottom: 1;
     border: round gray;
-}
-TabPane {
-    overflow-y: auto;
     padding: 1;
 }
 
-.row { height: auto; margin-bottom: 1;}
-.input_group { width: 1fr; margin-right: 2; height: auto; }
-.input_label { margin-top: 1; }
-#load_btn { width: auto; margin-top: 1; }
+#config_left_col {
+    width: 1fr;
+    padding-right: 2;
+    height: 100%;
+}
+
+#config_right_col {
+    width: 2fr;
+    height: 100%;
+}
+
+/* Ensure buttons take up consistent space and stack cleanly */
+#load_btn, #reload_btn {
+    width: 100%;
+    margin-top: 1;
+    margin-bottom: 0;
+}
+
+#config_display {
+    height: 1fr;
+    border: round $primary-muted;
+    background: $surface;
+}
+
 
-/* Bottom Section: Columns */
-#bottom_area { height: 1fr; margin-top: 1; }
-#control_panel { width: 6fr; padding-right: 2; }
+/* =========================================
+   Bottom Section: Columns
+   ========================================= */
+#bottom_area { height: 1fr; margin-top: 0; }
+#control_panel { width: 6fr; padding-right: 1; }
 #log_panel { width: 5fr; height: 1fr; border: solid white; }
 
 /* Controls, Dropdowns, and Status */
@@ -28,6 +49,16 @@ TabPane {
     margin-bottom: 1;
 }
 
+#log_header_row {
+    height: auto;
+    margin-bottom: 0;
+}
+
+#log_status_row {
+    height: auto;
+    margin-bottom: 1;
+}
+
 #scenario_select {
     width: 1fr;
     margin-right: 1;
@@ -35,23 +66,26 @@ TabPane {
 
 #pipeline_status_label {
     width: 1fr;
-    height: 100%;
+    height: 3;
     content-align: center middle;
     background: $boost;
     border: round $primary;
     text-style: bold;
 }
 
-Button { margin-top: 1; margin-bottom: 1; width: 100%; }
+Button { margin-top: 0; margin-bottom: 0; width: 100%; }
 #active_controls { height: auto; }
 #active_controls Button { width: 1fr; margin-right: 1; }
+#btn_start { margin-bottom: 1; }
 RichLog { height: 1fr; }
 
-/* Progress Bars Setup */
-#progress_area { height: auto; margin-top: 1; margin-bottom: 1; }
+/* =========================================
+   Progress Bars Setup
+   ========================================= */
+#progress_area { height: auto; margin-top: 0; margin-bottom: 0; }
 
 .progress_row {
-    height: auto;
+    height: 4;
     layout: horizontal;
     display: none;
     border: round $primary-muted;
@@ -61,7 +95,7 @@ RichLog { height: 1fr; }
 }
 
 .progress_title {
-    width: 25;
+    width: 16;
     height: auto;
     content-align: left middle;
     padding-right: 1;
@@ -70,11 +104,11 @@ RichLog { height: 1fr; }
     border-right: solid $primary-muted;
 }
 
-.progress_bar { width: 1fr; margin-top: 1; }
-ProgressBar ETAStatus { width: 9; }
+.progress_bar { width: 1fr; height: 2; margin-top: 0; }
+ProgressBar ETAStatus { width: auto; }
 
 .progress_stats {
-    width: 12;
+    width: 8;
     height: auto;
     content-align: right middle;
     padding-left: 1;
@@ -87,7 +121,9 @@ ProgressBar ETAStatus { width: 9; }
 #progress_bar_3 .bar--complete { background: $error; }
 #progress_bar_4 .bar--complete { background: $accent; }
 
-/* Modal Screen CSS */
+/* =========================================
+   Modal Screen CSS
+   ========================================= */
 OverwriteConfirmScreen { background: $background 80%; }
 #warning_label { width: 100%; height: auto; }
 #dialog { padding: 1 2; border: thick $error; width: 60; height: auto; background: $surface; }

util/tasks.py

@@ -1,85 +0,0 @@
-import time
-from typing import Any, Dict
-
-from util.progress import ProgressTracker
-from util.ui_logger import PipelineLogger
-
-
-def dummy_task(
-    tracker: ProgressTracker,
-    logger: PipelineLogger,
-    config: Dict[str, Any],
-    state: Dict[str, Any],
-):
-    stop_event = state["events"]["stop"]
-    pause_event = state["events"]["pause"]
-
-    steps = 5
-    # The title is now set gracefully by the parent injecting the sub_tracker,
-    # so we just initialize the total.
-    tracker.update(total=steps, advance=0)
-
-    logger.info("Initializing process...")
-
-    for i in range(steps):
-        if stop_event.is_set():
-            raise InterruptedError("Pipeline execution stopped by user.")
-
-        while pause_event.is_set():
-            time.sleep(0.5)
-            if stop_event.is_set():
-                raise InterruptedError(
-                    "Pipeline execution stopped by user while paused."
-                )
-
-        # Child Process (Sub Progress Tracker)
-        sub_steps = 10
-        sub_tracker = tracker.get_sub_tracker(f"Batch {i + 1}")
-        sub_tracker.update(total=sub_steps, advance=0)
-
-        for j in range(sub_steps):
-            if stop_event.is_set():
-                raise InterruptedError("Pipeline execution stopped by user.")
-            time.sleep(0.05)
-            sub_tracker.update(advance=1)  # Advance sub task
-
-        tracker.update(advance=1)  # Advance main task (clears sub task)
-
-        if i == 2:
-            logger.info("Halfway through current task execution...")
-
-
-PIPELINE_TASKS = {
-    "load_data": {"task_name": "Load Image and ADNIMERGE", "task_func": dummy_task},
-    "train_regular": {"task_name": "Train Regular Models", "task_func": dummy_task},
-    "train_bayesian": {"task_name": "Train Bayesian Models", "task_func": dummy_task},
-    "evaluate_regular": {
-        "task_name": "Evaluate Regular Models",
-        "task_func": dummy_task,
-    },
-    "evaluate_bayesian": {
-        "task_name": "Evaluate Bayesian Models",
-        "task_func": dummy_task,
-    },
-}
-
-SCENARIOS = {
-    "scen_train_all": {
-        "label": "1. Train, Evaluate, & Noise Analysis",
-        "tasks": [
-            "load_data",
-            "train_regular",
-            "train_bayesian",
-            "evaluate_regular",
-            "evaluate_bayesian",
-        ],
-    },
-    "scen_load_all": {
-        "label": "2. Load, Evaluate, & Noise Analysis",
-        "tasks": ["load_data", "evaluate_regular", "evaluate_bayesian"],
-    },
-    "scen_load_eval": {
-        "label": "3. Load & Evaluate (Skip Noise)",
-        "tasks": ["load_data", "evaluate_regular", "evaluate_bayesian"],
-    },
-}