{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":59093,"databundleVersionId":7469972,"sourceType":"competition"},{"sourceId":7899529,"sourceType":"datasetVersion","datasetId":4639130},{"sourceId":7908175,"sourceType":"datasetVersion","datasetId":4645532},{"sourceId":7977395,"sourceType":"datasetVersion","datasetId":4646663}],"dockerImageVersionId":30674,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"%pip install /kaggle/input/tmp63271321/einops-0.7.0-py3-none-any.whl","metadata":{"execution":{"iopub.status.busy":"2024-03-29T21:01:44.413420Z","iopub.execute_input":"2024-03-29T21:01:44.414125Z","iopub.status.idle":"2024-03-29T21:02:17.580785Z","shell.execute_reply.started":"2024-03-29T21:01:44.414096Z","shell.execute_reply":"2024-03-29T21:02:17.579675Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport einops\nimport cv2\nimport pandas as pd\nimport numpy as np\nfrom glob import glob\nimport matplotlib.pyplot as plt \nimport torch as t\nfrom torch.utils.data import Dataset, DataLoader\nimport torch as t\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom scipy.signal import butter, sosfilt, iirnotch, lfilter\n\ndevice = 'cuda' if t.cuda.is_available() else 'cpu'","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-03-29T21:02:17.582694Z","iopub.execute_input":"2024-03-29T21:02:17.583007Z","iopub.status.idle":"2024-03-29T21:02:23.557662Z","shell.execute_reply.started":"2024-03-29T21:02:17.582979Z","shell.execute_reply":"2024-03-29T21:02:23.556763Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_path = '/kaggle/input/hms-harmful-brain-activity-classification/test_eegs/'\nBASE_PATH = \"/kaggle/input/hms-harmful-brain-activity-classification\"\nclass_names = ['Seizure', 'LPD', 'GPD', 'LRDA','GRDA', 'Other']\nTARGETS = ['seizure_vote', 'lpd_vote', 'gpd_vote', 'lrda_vote', 'grda_vote', 'other_vote']\nFEATS_FOR_REAL = ['Fp1', 'F3', 'C3', 'P3', 'F7', 'T3', 'T5', 'O1', 'Fz', 'Cz', 'Pz', 'Fp2', 'F4', 'C4', 'P4', 'F8', 'T4', 'T6', 'O2', 'EKG']\nGROUPS_IDS = [\n    [0, 1, 2, 3, 7],\n    [0, 4, 5, 6, 7],\n    [11, 12, 13, 14, 18],\n    [11, 15, 16, 17, 18],\n]\n\nGROUPS_IDS2 = [\n    [0, 1, 2, 3, 7],\n    [0, 4, 5, 6, 7],\n    [11, 12, 13, 14, 18],\n    [11, 15, 16, 17, 18],\n    [8, 9, 10],\n]","metadata":{"execution":{"iopub.status.busy":"2024-03-29T21:02:23.558866Z","iopub.execute_input":"2024-03-29T21:02:23.559278Z","iopub.status.idle":"2024-03-29T21:02:23.566734Z","shell.execute_reply.started":"2024-03-29T21:02:23.559251Z","shell.execute_reply":"2024-03-29T21:02:23.565865Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_df = pd.read_csv(f'{BASE_PATH}/test.csv')","metadata":{"execution":{"iopub.status.busy":"2024-03-29T21:02:23.569416Z","iopub.execute_input":"2024-03-29T21:02:23.569757Z","iopub.status.idle":"2024-03-29T21:02:23.591907Z","shell.execute_reply.started":"2024-03-29T21:02:23.569727Z","shell.execute_reply":"2024-03-29T21:02:23.591058Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def clip(data, bound=300):\n    return np.clip(data, -bound, bound)\n\ndef robust_norm(data):\n    median = np.median(data, axis=0)\n    q75, q25 = np.percentile(data, [75 ,25], axis=0)\n    iqr = q75 - q25\n    iqr[iqr < 1e-6] = 1e-6 # numerical stability\n    return (data - median) / iqr\n\ndef band_filter(data, low=1, high=70, fs=200, order=4):\n    sos = butter(N=order, Wn=[low, high], btype='bandpass', fs=fs, output='sos')\n    return sosfilt(sos, data, axis=0)\n\ndef notch_filter(data, notch_freq=60, Q=30, fs=200):\n    b_notch, a_notch = iirnotch(notch_freq, Q, fs)\n    return lfilter(b_notch, a_notch, data, axis=0)\n\ndef filters(data):\n    data = band_filter(data)\n    data = notch_filter(data)\n#     data = robust_norm(data)\n    return data\n\nclass EEGTestDataset(Dataset):\n    def __init__(self):\n        super().__init__()\n        self.dataframe = test_df\n\n    def __len__(self):\n        return len(self.dataframe)\n\n    def __getitem__(self, idx):\n        row = self.dataframe.iloc[idx]\n        eeg_id = row['eeg_id']\n        parq_path = f'{test_path}{eeg_id}.parquet'\n        eeg = pd.read_parquet(parq_path)\n        eeg = eeg.ffill(axis=0)\n        eeg = eeg.fillna(0)\n        filtered_eeg = filters(eeg[FEATS_FOR_REAL].values)\n        samples = t.tensor(filtered_eeg).float()        \n        return samples","metadata":{"execution":{"iopub.status.busy":"2024-03-29T21:02:23.593068Z","iopub.execute_input":"2024-03-29T21:02:23.593376Z","iopub.status.idle":"2024-03-29T21:02:23.605674Z","shell.execute_reply.started":"2024-03-29T21:02:23.593333Z","shell.execute_reply":"2024-03-29T21:02:23.604792Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"batch_size = 150\ntest_dataset = EEGTestDataset()\ntest_dataloader = DataLoader(test_dataset, batch_size=batch_size, num_workers=4, prefetch_factor=2)","metadata":{"execution":{"iopub.status.busy":"2024-03-29T21:02:23.606616Z","iopub.execute_input":"2024-03-29T21:02:23.606954Z","iopub.status.idle":"2024-03-29T21:02:23.622799Z","shell.execute_reply.started":"2024-03-29T21:02:23.606929Z","shell.execute_reply":"2024-03-29T21:02:23.621859Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# class ConvBlock(nn.Module):\n#     def __init__(self, d_in, d_out, kernel_size, drop):\n#         super().__init__()\n#         self.model = nn.Sequential(\n#             nn.Conv1d(d_in, d_out, kernel_size=kernel_size, padding='same', stride=1),\n#             nn.InstanceNorm1d(d_out),\n#             nn.ReLU(),\n#             nn.Dropout(drop),\n#             nn.Conv1d(d_out, d_out, kernel_size=kernel_size, padding='same', stride=1),\n#             nn.InstanceNorm1d(d_out),\n#             nn.ReLU(),\n#             nn.Dropout(drop),\n#             nn.Conv1d(d_out, d_out, kernel_size=kernel_size, padding='same', stride=1),\n#             nn.InstanceNorm1d(d_out),\n#             nn.ReLU(),\n#             nn.Dropout(drop),\n#             nn.MaxPool1d(kernel_size=2, stride=2, padding=0), # reduce sequence size by 2\n#         )\n#     def forward(self, x):\n#         # TODO: add skip for training speed\n#         return self.model(x)\n        \n# class Model(nn.Module):\n#     def __init__(self, in_channels=4, gru_hidden_size=128, drop=0.5):\n#         super().__init__()\n#         self.d_split = len(GROUPS_IDS)\n#         self.pre_out = in_channels * 4\n#         self.gru_hidden_size = gru_hidden_size\n        \n#         self.pre_process = nn.Sequential(\n#             nn.InstanceNorm1d(in_channels),\n#             # nn.LayerNorm(normalized_shape=[in_channels, 10000]),\n#             # nn.BatchNorm1d(in_channels, momentum=None),\n#             # use conv1d as a denoiser\n#             # -- block 1 --\n#             ConvBlock(in_channels, in_channels * 2, kernel_size=3, drop=drop),\n#             # nn.BatchNorm1d(in_channels * 2, momentum=None),\n#             # nn.LayerNorm(normalized_shape=[in_channels * 2, 5000]),\n#             # -- block 2 --\n#             ConvBlock(in_channels * 2, in_channels * 4, kernel_size=5, drop=drop),\n#             # nn.BatchNorm1d(in_channels * 4, momentum=None),\n#             # nn.LayerNorm(normalized_shape=[in_channels * 4, 2500]),\n#             # -- block 3 --\n#             ConvBlock(in_channels * 4, in_channels * 4, kernel_size=7, drop=drop),\n#             # nn.BatchNorm1d(in_channels * 4, momentum=None), # re-enable one more to force training ?\n#             # nn.LayerNorm(normalized_shape=[in_channels * 4, 1250]),\n#         )\n        \n#         # TODO: add a learnable first state for GRU or check what is the default\n#         self.gru = nn.GRU(self.pre_out, self.gru_hidden_size, num_layers=1, batch_first=True, bidirectional=True)\n\n#         self.post_gru = nn.Sequential(\n#             # nn.BatchNorm1d(self.gru_hidden_size * 2, momentum=None),\n#             nn.LayerNorm(normalized_shape=[self.gru_hidden_size * 2]),\n#             nn.Dropout(drop),\n#             nn.Linear(self.gru_hidden_size * 2, self.gru_hidden_size),\n#             nn.ReLU(),\n#             nn.Dropout(drop),\n#             nn.Linear(self.gru_hidden_size, self.gru_hidden_size),\n#         )\n\n#         self.head = nn.Sequential(\n#             # nn.BatchNorm1d(self.gru_hidden_size * self.d_split, momentum=None),\n#             nn.LayerNorm(normalized_shape=[self.gru_hidden_size * self.d_split]),\n#             nn.Linear(self.gru_hidden_size * self.d_split, self.gru_hidden_size * 2),\n#             nn.ReLU(),\n#             nn.Dropout(drop),\n#             nn.Linear(self.gru_hidden_size * 2, 6)\n#         )\n\n#     def montage(self, x):\n#         splits = [x[:, :, group] for group in GROUPS_IDS]\n#         splits = [s[:, :, :-1] - s[:, :, 1:] for s in splits]\n#         return einops.rearrange(t.stack(splits, dim=0), 'group batch seq channel -> (group batch) seq channel')\n\n#     def forward(self, x: ('batch', 'seq', 'channel')):\n#         # separate the input into 4 montages (LP, LL, RP, RR)\n#         x = self.montage(x)\n#         # pre_process: (batch, channel, seq) → (batch / 4, channel * 4, seq)\n#         x = x.permute((0, 2, 1))\n#         x = self.pre_process(x)\n#         x = x.permute((0, 2, 1))\n#         # GRU: (batch, seq, input_size), [(2 * num_layers, batch, hidden_size)] → (batch, seq, 2 * hidden_size)\n#         x, _ = self.gru(x)\n#         x = x[:, -1, :]\n#         # MLP post GRU\n#         x = self.post_gru(x)\n#         # unfold the splits\n#         x = einops.rearrange(x, '(group batch) hidden -> batch (hidden group)', group=self.d_split)\n#         # head: (batch, 2 * hidden_size) → (batch, 6)\n#         x = self.head(x)\n#         # out: → (batch, 6)\n#         return x","metadata":{"execution":{"iopub.status.busy":"2024-03-29T21:02:23.623962Z","iopub.execute_input":"2024-03-29T21:02:23.624223Z","iopub.status.idle":"2024-03-29T21:02:23.637088Z","shell.execute_reply.started":"2024-03-29T21:02:23.624201Z","shell.execute_reply":"2024-03-29T21:02:23.636160Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ConvBlock(nn.Module):\n    def __init__(self, d_in, d_out, kernel_size, drop):\n        super().__init__()\n        self.model = nn.Sequential(\n            nn.Conv1d(d_in, d_out, kernel_size=kernel_size, padding='same', stride=1),\n            nn.InstanceNorm1d(d_out),\n            nn.ReLU(),\n            nn.Dropout(drop),\n            nn.Conv1d(d_out, d_out, kernel_size=kernel_size, padding='same', stride=1),\n            # nn.InstanceNorm1d(d_out),\n            nn.ReLU(),\n            nn.Dropout(drop),\n            nn.Conv1d(d_out, d_out, kernel_size=kernel_size, padding='same', stride=1),\n            # nn.InstanceNorm1d(d_out),\n            nn.ReLU(),\n            nn.Dropout(drop),\n            nn.MaxPool1d(kernel_size=2, stride=2, padding=0), # reduce sequence size by 2\n        )\n    def forward(self, x):\n        # TODO: add skip for training speed\n        return self.model(x)\n        \nclass Model(nn.Module):\n    def __init__(self, in_channels=19, gru_hidden_size=128, drop=0.5):\n    # def __init__(self, in_channels=4, gru_hidden_size=256, drop=0.5):\n        super().__init__()\n        # TODO double the size of channels, see if it helps?\n        in_channels *= 2\n        self.d_split = len(GROUPS_IDS)\n        self.pre_out = in_channels * 4\n        self.gru_hidden_size = gru_hidden_size\n        \n        self.pre_process = nn.Sequential(\n            # nn.InstanceNorm1d(in_channels),\n            # nn.LayerNorm(normalized_shape=[in_channels, 10000]),\n            # nn.BatchNorm1d(in_channels, momentum=None),\n            # use conv1d as a denoiser\n            # -- block 1 --\n            ConvBlock(in_channels // 2, in_channels * 2, kernel_size=3, drop=drop),\n            # nn.BatchNorm1d(in_channels * 2, momentum=None),\n            # nn.LayerNorm(normalized_shape=[in_channels * 2, 5000]),\n            # -- block 2 --\n            ConvBlock(in_channels * 2, in_channels * 4, kernel_size=5, drop=drop),\n            # nn.BatchNorm1d(in_channels * 4, momentum=None),\n            # nn.LayerNorm(normalized_shape=[in_channels * 4, 2500]),\n            # -- block 3 --\n            ConvBlock(in_channels * 4, in_channels * 4, kernel_size=7, drop=drop),\n            # nn.BatchNorm1d(in_channels * 4, momentum=None), # re-enable one more to force training ?\n            # nn.LayerNorm(normalized_shape=[in_channels * 4, 1250]),\n        )\n        \n        # TODO: add a learnable first state for GRU or check what is the default\n        self.gru = nn.GRU(self.pre_out, self.gru_hidden_size, num_layers=1, batch_first=True, bidirectional=True)\n\n        self.post_gru = nn.Sequential(\n            # nn.BatchNorm1d(self.gru_hidden_size * 2, momentum=None),\n            nn.LayerNorm(normalized_shape=[self.gru_hidden_size * 2]),\n            nn.Dropout(drop),\n            # nn.LayerNorm(normalized_shape=self.gru_hidden_size * 2),\n            nn.Linear(self.gru_hidden_size * 2, self.gru_hidden_size),\n            nn.ReLU(),\n            nn.Dropout(drop),\n            # nn.LayerNorm(normalized_shape=self.gru_hidden_size),\n            nn.Linear(self.gru_hidden_size, self.gru_hidden_size),\n            nn.ReLU(),\n            nn.Dropout(drop),\n            # nn.LayerNorm(normalized_shape=self.gru_hidden_size),\n            nn.Linear(self.gru_hidden_size, 6),\n        )\n\n        # self.head = nn.Sequential(\n        #     # nn.BatchNorm1d(self.gru_hidden_size * self.d_split, momentum=None),\n        #     # nn.LayerNorm(normalized_shape=[self.gru_hidden_size * self.d_split]),\n        #     nn.Linear(self.gru_hidden_size * self.d_split, self.gru_hidden_size * 2),\n        #     nn.ReLU(),\n        #     nn.Dropout(drop),\n        #     nn.Linear(self.gru_hidden_size * 2, 6)\n        # )\n\n    def montage(self, x):\n        splits = [x[:, :, group] for group in GROUPS_IDS]\n        splits = [s[:, :, :-1] - s[:, :, 1:] for s in splits]\n        return einops.rearrange(t.stack(splits, dim=0), 'group batch seq channel -> (group batch) seq channel')\n    \n    def montage2(self, x):\n        splits = [x[:, :, group] for group in GROUPS_IDS2]\n        splits = [s[:, :, :-1] - s[:, :, 1:] for s in splits]\n        splits.append(x[:, :, 19:20]) # add EKG\n        return t.cat(splits, dim=2)\n        # return einops.rearrange(t.cat(splits, dim=2), 'group batch seq channel -> group seq (batch channel)')\n\n    def forward(self, x: ('batch', 'seq', 'channel')):\n        # separate the input into 4 montages (LP, LL, RP, RR)\n        x = self.montage2(x)\n        # pre_process: (batch, channel, seq) → (batch / 4, channel * 4, seq)\n        x = x.permute((0, 2, 1))\n        x = self.pre_process(x)\n        x = x.permute((0, 2, 1))\n        # GRU: (batch, seq, input_size), [(2 * num_layers, batch, hidden_size)] → (batch, seq, 2 * hidden_size)\n        x, _ = self.gru(x)\n        x = x[:, -1, :]\n        # MLP post GRU\n        x = self.post_gru(x)\n        # unfold the splits\n        # x = einops.rearrange(x, '(group batch) hidden -> batch (hidden group)', group=self.d_split)\n        # head: (batch, 2 * hidden_size) → (batch, 6)\n        # x = self.head(x)\n        # out: → (batch, 6)\n        return x","metadata":{"execution":{"iopub.status.busy":"2024-03-29T21:02:23.638340Z","iopub.execute_input":"2024-03-29T21:02:23.638581Z","iopub.status.idle":"2024-03-29T21:02:23.660488Z","shell.execute_reply.started":"2024-03-29T21:02:23.638560Z","shell.execute_reply":"2024-03-29T21:02:23.659552Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = Model().to(device)\nmodel.load_state_dict(t.load('/kaggle/input/tmp63271321/gru-4-splits_2024-03-29_21h25.pt', map_location=device))","metadata":{"execution":{"iopub.status.busy":"2024-03-29T21:02:23.661409Z","iopub.execute_input":"2024-03-29T21:02:23.661652Z","iopub.status.idle":"2024-03-29T21:02:24.221651Z","shell.execute_reply.started":"2024-03-29T21:02:23.661630Z","shell.execute_reply":"2024-03-29T21:02:24.220698Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def round_zero(res, threshold=1e-3):\n    res[res < threshold] = 0\n    res = res / res.sum(dim=-1, keepdim=True)\n    return res\n\nres = []\nwith t.no_grad():\n    model.eval()\n\n    for batch in test_dataloader:\n        batch = batch.to(device)\n        prob = model(batch).softmax(-1)\n        res.append(prob.detach().cpu())\n\n    res = t.cat(res, dim=0)\n#     res = round_zero(res)","metadata":{"execution":{"iopub.status.busy":"2024-03-29T21:02:24.224475Z","iopub.execute_input":"2024-03-29T21:02:24.224766Z","iopub.status.idle":"2024-03-29T21:02:25.279718Z","shell.execute_reply.started":"2024-03-29T21:02:24.224741Z","shell.execute_reply":"2024-03-29T21:02:25.278598Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sub = test_df[[\"eeg_id\"]].copy()\nsub[TARGETS] = res\nsub.to_csv('submission.csv',index=False)\n\nprint('Submission shape',sub.shape)\nsub.head()","metadata":{"execution":{"iopub.status.busy":"2024-03-29T21:02:25.281174Z","iopub.execute_input":"2024-03-29T21:02:25.281484Z","iopub.status.idle":"2024-03-29T21:02:25.308518Z","shell.execute_reply.started":"2024-03-29T21:02:25.281457Z","shell.execute_reply":"2024-03-29T21:02:25.307592Z"},"trusted":true},"execution_count":null,"outputs":[]}]}