diff --git a/data/src/consts.py b/data/src/consts.py index 25481fa..f3884eb 100644 --- a/data/src/consts.py +++ b/data/src/consts.py @@ -1,7 +1,8 @@ import numpy as np SAMPLING_RATE = 250 -SAMPLES_PER_5_SEC = 1250 +WINDOW_TIME_FRAME = 5 +SAMPLES_PER_5_SEC = SAMPLING_RATE * WINDOW_TIME_FRAME ADAPT_TRIALS = 8 EXPECTED_FREQS = [6.66, 7.50, 8.57, 10.00, 12.00] diff --git a/data/src/data_reader.py b/data/src/data_reader.py index a68e095..7420e24 100644 --- a/data/src/data_reader.py +++ b/data/src/data_reader.py @@ -1,95 +1,130 @@ from collections import defaultdict import scipy.io from data_classes import * -from consts import * - - -def load_data(filename: str) -> tuple[np.ndarray, np.ndarray, list[int]]: - mat_data = scipy.io.loadmat(filename) - eeg_data = mat_data[EEG_KEY] - din_data = mat_data[DIN_KEY] - - latencies_ms = [int(din_data[1, i].item()) for i in range(din_data.shape[1])] - latencies = [int(round(l * SAMPLING_RATE / 1000)) for l in latencies_ms] - - return eeg_data, latencies_ms, latencies - - -def group_din_markers_into_trials( - latencies: list[int], gap_threshold_samples: int -) -> list[trial_group]: - trial_groups = [] - current_start_idx = 0 - current_group_start = latencies[0] - current_group_end = latencies[0] - - for i in range(1, len(latencies)): - if latencies[i] - latencies[i - 1] > gap_threshold_samples: - trial_groups.append( - trial_group( - start_sample=current_group_start, - end_sample=current_group_end, - start_idx=current_start_idx, - end_idx=i - 1, +from consts import ( + EEG_KEY, + DIN_KEY, + SAMPLING_RATE, + EXPECTED_FREQS, + ADAPT_TRIALS, + WINDOW_TIME_FRAME, +) +from abc import ABC, abstractmethod + + +class DataReader(ABC): + @abstractmethod + def _load_data(self, filename: str) -> tuple[np.ndarray, np.ndarray, list[int]]: + pass + + @abstractmethod + def get_all_trials(self, filename: str) -> list[trial_info]: + pass + + +class MatDataReader(DataReader): + def _load_data( + self, filename: str, sampling_rate: int = SAMPLING_RATE + ) -> tuple[np.ndarray, np.ndarray, list[int]]: + mat_data = scipy.io.loadmat(filename) + eeg_data = mat_data[EEG_KEY] + din_data = mat_data[DIN_KEY] + + latencies_ms = [int(din_data[1, i].item()) for i in range(din_data.shape[1])] + latencies = [int(round(l * sampling_rate / 1000)) for l in latencies_ms] + + return eeg_data, latencies_ms, latencies + + def _group_din_markers_into_trials( + self, latencies: list[int], gap_threshold_samples: int + ) -> list[trial_group]: + trial_groups = [] + current_start_idx = 0 + current_group_start = latencies[0] + current_group_end = latencies[0] + + for i in range(1, len(latencies)): + if latencies[i] - latencies[i - 1] > gap_threshold_samples: + trial_groups.append( + trial_group( + start_sample=current_group_start, + end_sample=current_group_end, + start_idx=current_start_idx, + end_idx=i - 1, + ) ) + current_group_start = latencies[i] + current_start_idx = i + current_group_end = latencies[i] + + trial_groups.append( + trial_group( + start_sample=current_group_start, + end_sample=current_group_end, + start_idx=current_start_idx, + end_idx=len(latencies) - 1, ) - current_group_start = latencies[i] - current_start_idx = i - current_group_end = latencies[i] - - trial_groups.append( - trial_group( - start_sample=current_group_start, - end_sample=current_group_end, - start_idx=current_start_idx, - end_idx=len(latencies) - 1, ) - ) - - return trial_groups + return trial_groups + + def _extract_trials( + self, + eeg_data: np.ndarray, + latencies_ms: list[int], + trial_groups: list[trial_group], + adapt_trials: int = ADAPT_TRIALS, + trial_time_frame_sec: int = WINDOW_TIME_FRAME, + sampling_rate: int = SAMPLING_RATE, + ) -> list[trial_info]: + all_trials = [] + samples_per_trial = sampling_rate * trial_time_frame_sec + + for i, tg in enumerate(trial_groups): + start_sample = tg.start_sample + epoch = eeg_data[:, start_sample : start_sample + samples_per_trial] + + if epoch.shape[1] != samples_per_trial: + continue + + dins_ms = latencies_ms[tg.start_idx : tg.end_idx + 1] + n_dins = len(dins_ms) + true_freq = 0.0 + closest = 0.0 + + if n_dins > 1: + intervals = [ + dins_ms[j] - dins_ms[j - 1] for j in range(1, len(dins_ms)) + ] + true_freq = (1000.0 / np.mean(intervals)) / 2.0 + closest = min(EXPECTED_FREQS, key=lambda f: abs(f - true_freq)) + + trial_type = TrialType.ADAPT if i < adapt_trials else TrialType.TEST + + trial_record = trial_info( + epoch=epoch, + trial=i, + true_freq=round(true_freq, 2), + closest_freq=closest, + n_dins=n_dins, + type=trial_type, + ) -def extract_trials( - eeg_data: np.ndarray, latencies_ms: list[int], trial_groups: list[trial_group] -) -> list[trial_info]: - all_trials = [] - - for i, tg in enumerate(trial_groups): - start_sample = tg.start_sample - epoch = eeg_data[:, start_sample : start_sample + SAMPLES_PER_5_SEC] - - if epoch.shape[1] != SAMPLES_PER_5_SEC: - continue - - dins_ms = latencies_ms[tg.start_idx : tg.end_idx + 1] - n_dins = len(dins_ms) - true_freq = 0.0 - closest = 0.0 - - if n_dins > 1: - intervals = [dins_ms[j] - dins_ms[j - 1] for j in range(1, len(dins_ms))] - true_freq = (1000.0 / np.mean(intervals)) / 2.0 - closest = min(EXPECTED_FREQS, key=lambda f: abs(f - true_freq)) - - trial_type = TrialType.ADAPT if i < ADAPT_TRIALS else TrialType.TEST + all_trials.append(trial_record) - trial_record = trial_info( - epoch=epoch, - trial=i, - true_freq=round(true_freq, 2), - closest_freq=closest, - n_dins=n_dins, - type=trial_type, - ) + return all_trials - all_trials.append(trial_record) + def get_all_trials(self, filename: str) -> list[trial_info]: + eeg_data, latencies_ms, latencies = self._load_data(filename) + trial_groups = self._group_din_markers_into_trials(latencies, 250) + all_trials = self._extract_trials(eeg_data, latencies_ms, trial_groups) - return all_trials + return all_trials -def main(filename: str) -> list[trial_info]: - eeg_data, latencies_ms, latencies = load_data(filename) - trial_groups = group_din_markers_into_trials(latencies, 250) - all_trials = extract_trials(eeg_data, latencies_ms, trial_groups) +class CSVDataReader(DataReader): + def _load_data(self, filename): + return super()._load_data(filename) - return all_trials + def get_all_trials(self, filename): + return super().get_all_trials(filename) diff --git a/data/src/fft_analysis.py b/data/src/fft_analysis.py index f8b11e5..f1524c4 100644 --- a/data/src/fft_analysis.py +++ b/data/src/fft_analysis.py @@ -4,7 +4,7 @@ from consts import ( SAMPLING_RATE, - SAMPLES_PER_5_SEC, + WINDOW_TIME_FRAME, SNR_NEIGHBOR_BINS, SNR_EXCLUDE_BINS, SNR_THRESHOLD_DB, @@ -16,8 +16,8 @@ from data_classes import trial_info, TrialType from utils import group_trials_by_frequency, get_trials_by_type -N = SAMPLES_PER_5_SEC -freqs_fft = np.fft.rfftfreq(N, d=1.0 / SAMPLING_RATE) +samples_per_trial = SAMPLING_RATE * WINDOW_TIME_FRAME +freqs_fft = np.fft.rfftfreq(samples_per_trial, d=1.0 / SAMPLING_RATE) def compute_rfft(signal: np.ndarray) -> np.ndarray: @@ -27,7 +27,7 @@ def compute_rfft(signal: np.ndarray) -> np.ndarray: def compute_fft_magnitude(signal: np.ndarray) -> np.ndarray: fft_result = compute_rfft(signal) - magnitude = (2.0 / N) * np.abs(fft_result) + magnitude = (2.0 / samples_per_trial) * np.abs(fft_result) return magnitude diff --git a/data/src/main.py b/data/src/main.py index 1202312..dd5fe96 100644 --- a/data/src/main.py +++ b/data/src/main.py @@ -1,5 +1,5 @@ import argparse -from data_reader import main as read_data +from data_reader import DataReader, MatDataReader, CSVDataReader from filter import main as filter_data from fft_analysis import main as fft_analysis_main from filter_plotter import ( @@ -14,10 +14,22 @@ from consts import OUTPUT_DIR import numpy as np +# Data file type to reader map +DATA_READERS = { + "csv": CSVDataReader, + "mat": MatDataReader, +} -def main(filename: str, plot_filter: bool, plot_fft: bool): + +def main( + filename: str, + plot_filter: bool, + plot_fft: bool, + reader: DataReader, + output_dir: str = OUTPUT_DIR, +): print("\nReading data...") - trials = read_data(filename) + trials = reader.get_all_trials(filename) # Save raw (unfiltered) epochs before applying any filtering raw_epochs = np.array([t.epoch for t in trials]) @@ -50,20 +62,32 @@ def main(filename: str, plot_filter: bool, plot_fft: bool): r["ch_idx"], r["channel"], r["snr_db"], - OUTPUT_DIR, + output_dir, ) for ef in all_results: - plot_grouped_fft(ef, all_results[ef], stats, OUTPUT_DIR) + plot_grouped_fft(ef, all_results[ef], stats, output_dir) - plot_snr_heatmap(stats, OUTPUT_DIR) + plot_snr_heatmap(stats, output_dir) print("\nPipeline complete!") if __name__ == "__main__": parser = argparse.ArgumentParser(description="Run the full SSVEP pipeline.") - parser.add_argument("--filename", type=str, help="Path to the .mat file") + parser.add_argument( + "--filename", + type=str, + required=True, + help="Path to the data file", + ) + parser.add_argument( + "--data_type", + type=str, + required=True, + choices=["csv", "mat"], + help="csv or mat", + ) parser.add_argument( "--plot_filter", action="store_true", @@ -74,6 +98,11 @@ def main(filename: str, plot_filter: bool, plot_fft: bool): action="store_true", help="Generate plots for the FFT results (SNR heatmap, FFT spectrum)", ) + parser.add_argument( + "--output", default=OUTPUT_DIR, type=str, help="output directory" + ) args = parser.parse_args() - main(args.filename, args.plot_filter, args.plot_fft) + reader_class = DATA_READERS.get(args.data_type) + + main(args.filename, args.plot_filter, args.plot_fft, reader_class(), args.output)