{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.14","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":59093,"databundleVersionId":7469972,"sourceType":"competition"}],"dockerImageVersionId":30786,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport mne\nfrom pathlib import Path\nfrom scipy.signal import butter, filtfilt\nfrom mne.preprocessing import ICA\nimport shutil\n\nclass EEGProcessor:\n    def __init__(self, input_path, sample_rate=200):\n        self.input_path = Path(input_path)\n        self.sample_rate = sample_rate\n        self.df = None\n        self.raw = None\n        self.filtered_raw = None\n        self.cleaned_raw = None\n        \n        # Change output to working directory\n        self.base_dir = Path(\"/kaggle/working\")\n        self.cleaned_dir = self.base_dir / \"cleaned_EEGs\"\n        self.unprocessed_dir = self.base_dir / \"unprocessed_EEGs\"\n        self.cleaned_dir.mkdir(exist_ok=True)\n        self.unprocessed_dir.mkdir(exist_ok=True)\n        \n        self.cleaned_path = self.cleaned_dir / self.input_path.name\n        self.unprocessed_path = self.unprocessed_dir / self.input_path.name\n\n    def butter_highpass_filter(self, data, cutoff, order=4):\n        try:\n            nyquist = 0.5 * self.sample_rate\n            normal_cutoff = cutoff / nyquist\n            b, a = butter(order, normal_cutoff, btype='high', analog=False)\n            return filtfilt(b, a, data, axis=-1)\n        except Exception:\n            return data\n\n    def load_data(self):\n        try:\n            self.df = pd.read_parquet(self.input_path)\n            if 'EKG' in self.df.columns:\n                self.df = self.df.drop('EKG', axis=1)\n            \n            data = self.df.to_numpy().T\n            ch_names = list(self.df.columns)\n            ch_types = ['eeg'] * len(ch_names)\n            info = mne.create_info(ch_names=ch_names, sfreq=self.sample_rate, ch_types=ch_types)\n            self.raw = mne.io.RawArray(data, info)\n            return self\n        except Exception as e:\n            print(f\"⚠️ Load failed\")\n            raise\n\n    def apply_preprocessing(self, highpass_cutoff=1.0):\n        try:\n            data = self.raw.get_data()\n            filtered_data = np.zeros_like(data)\n            for i in range(len(self.raw.ch_names)):\n                filtered_data[i] = self.butter_highpass_filter(\n                    data[i], cutoff=highpass_cutoff\n                )\n            \n            self.filtered_raw = mne.io.RawArray(filtered_data, self.raw.info)\n            return self\n        except Exception:\n            self.filtered_raw = self.raw.copy()\n            return self\n\n    def apply_automatic_ica(self, n_components=15):\n        try:\n            if self.filtered_raw is None:\n                self.filtered_raw = self.raw.copy()\n                \n            self.ica = ICA(n_components=n_components, random_state=97, max_iter=1000)\n            self.ica.fit(self.filtered_raw)\n            \n            components = self.ica.get_sources(self.filtered_raw).get_data()\n            artifacts_to_remove = []\n            \n            for idx, component in enumerate(components):\n                try:\n                    if self._is_eye_artifact(component, idx) or \\\n                       self._is_muscle_artifact(component) or \\\n                       self._is_line_noise(component):\n                        artifacts_to_remove.append(idx)\n                except Exception:\n                    continue\n\n            self.ica.exclude = artifacts_to_remove\n            self.cleaned_raw = self.filtered_raw.copy()\n            self.ica.apply(self.cleaned_raw)\n            \n            print(f\"↪ Removed {len(artifacts_to_remove)} artifacts\")\n            return self\n        except Exception:\n            self.cleaned_raw = self.filtered_raw.copy()\n            return self\n\n    def _is_eye_artifact(self, component, idx):\n        try:\n            weights = self.ica.get_components()[:, idx]\n            \n            frontal_channels = [i for i, ch in enumerate(self.filtered_raw.ch_names) \n                              if ch.startswith(('Fp', 'F'))]\n            if len(frontal_channels) > 0:\n                frontal_weight = np.mean(np.abs(weights[frontal_channels]))\n                other_weight = np.mean(np.abs(weights[~np.isin(np.arange(len(weights)), frontal_channels)]))\n                \n                if frontal_weight > 2 * other_weight:\n                    return True\n            \n            if np.percentile(np.abs(np.diff(component)), 99) > 4 * np.std(component):\n                return True\n                \n            return False\n        except Exception:\n            return False\n\n    def _is_muscle_artifact(self, component):\n        try:\n            freqs, psd = self._compute_psd(component)\n            \n            if len(freqs) == 0 or len(psd) == 0:\n                return False\n                \n            low_freq_power = np.mean(psd[(freqs >= 1) & (freqs <= 20)])\n            high_freq_power = np.mean(psd[(freqs >= 20) & (freqs <= 100)])\n            \n            if high_freq_power > 0.5 * low_freq_power:\n                return True\n                \n            return False\n        except Exception:\n            return False\n\n    def _is_line_noise(self, component):\n        try:\n            freqs, psd = self._compute_psd(component)\n            \n            if len(freqs) == 0 or len(psd) == 0:\n                return False\n                \n            for line_freq in [50, 60]:\n                mask = (freqs >= line_freq - 1) & (freqs <= line_freq + 1)\n                if np.any(mask) and np.max(psd[mask]) > 5 * np.median(psd):\n                    return True\n                    \n            return False\n        except Exception:\n            return False\n\n    def _compute_psd(self, component):\n        try:\n            return mne.time_frequency.psd_array_welch(\n                component,\n                sfreq=self.sample_rate,\n                fmin=1,\n                fmax=100,\n                n_fft=int(2 * self.sample_rate)\n            )\n        except Exception:\n            return np.array([]), np.array([])\n\n    def save_cleaned_data(self):\n        try:\n            if self.cleaned_raw is not None:\n                cleaned_data = self.cleaned_raw.get_data().T\n                cleaned_df = pd.DataFrame(cleaned_data, columns=self.cleaned_raw.ch_names)\n                cleaned_df.to_parquet(self.cleaned_path)\n                print(\"✓ Saved\")\n            else:\n                shutil.copy2(self.input_path, self.unprocessed_path)\n                print(\"⚠️ Copied to unprocessed\")\n            return self\n        except Exception:\n            shutil.copy2(self.input_path, self.unprocessed_path)\n            print(\"⚠️ Copied to unprocessed\")\n            return self\n\ndef process_multiple_recordings(input_directory, start_index=0, end_index=None):\n    input_dir = Path(input_directory)\n    parquet_files = list(input_dir.glob(\"*.parquet\"))\n    total_files = len(parquet_files)\n    \n    # Slice the list of files based on start and end indices\n    files_to_process = parquet_files[start_index:end_index]\n    selected_files_count = len(files_to_process)\n    \n    print(f\"Found {total_files} files. Processing {selected_files_count} files from index {start_index} to {end_index or total_files}.\")\n    \n    for i, file_path in enumerate(files_to_process, start=start_index + 1):\n        file_name = file_path.name\n        print(f\"\\n[{i}/{total_files}] {file_name}\")\n        try:\n            processor = EEGProcessor(file_path, sample_rate=200)\n            processor.load_data() \\\n                    .apply_preprocessing() \\\n                    .apply_automatic_ica() \\\n                    .save_cleaned_data()\n        except Exception:\n            try:\n                processor = EEGProcessor(file_path)\n                shutil.copy2(file_path, processor.unprocessed_path)\n                print(\"⚠️ Copied to unprocessed\")\n            except Exception:\n                print(\"❌ Failed completely\")\n\n\nif __name__ == \"__main__\":\n    process_multiple_recordings(\n        \"/kaggle/input/hms-harmful-brain-activity-classification/train_eegs\",\n        start_index=8000,\n        end_index=10000,\n    )\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-11-16T14:17:25.405376Z","iopub.execute_input":"2024-11-16T14:17:25.405853Z","iopub.status.idle":"2024-11-16T14:17:44.249816Z","shell.execute_reply.started":"2024-11-16T14:17:25.405802Z","shell.execute_reply":"2024-11-16T14:17:44.248573Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{},"outputs":[],"execution_count":null}]}