{"metadata":{"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":87230,"databundleVersionId":9911932,"sourceType":"competition"}],"dockerImageVersionId":30786,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true},"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.10.14"},"papermill":{"default_parameters":{},"duration":4943.73701,"end_time":"2024-10-29T10:41:54.878842","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2024-10-29T09:19:31.141832","version":"2.6.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install einops\n!pip install monai==1.3.2","metadata":{"papermill":{"duration":25.753614,"end_time":"2024-10-29T09:19:59.601090","exception":false,"start_time":"2024-10-29T09:19:33.847476","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-11-15T11:27:55.354074Z","iopub.execute_input":"2024-11-15T11:27:55.354429Z","iopub.status.idle":"2024-11-15T11:28:21.644558Z","shell.execute_reply.started":"2024-11-15T11:27:55.354393Z","shell.execute_reply":"2024-11-15T11:28:21.643585Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nimport os\nimport glob\nimport matplotlib.pyplot as plt\nfrom sklearn.preprocessing import StandardScaler\nfrom torch.utils.data import Dataset, DataLoader, Subset, TensorDataset\nimport pandas as pd\nfrom sklearn.impute import SimpleImputer\nimport torch\nfrom sklearn.model_selection import KFold\nfrom monai.networks.nets import SwinUNETR, DynUNet, UNet\nfrom torch import nn\nfrom transformers import get_cosine_schedule_with_warmup\nfrom sklearn.metrics import mean_absolute_error\nimport torch.optim as optim\nfrom tqdm import tqdm\nfrom matplotlib import colors\nimport gc","metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","papermill":{"duration":45.222998,"end_time":"2024-10-29T09:20:44.830452","exception":false,"start_time":"2024-10-29T09:19:59.607454","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-11-15T11:28:21.646842Z","iopub.execute_input":"2024-11-15T11:28:21.647266Z","iopub.status.idle":"2024-11-15T11:29:07.242487Z","shell.execute_reply.started":"2024-11-15T11:28:21.647218Z","shell.execute_reply":"2024-11-15T11:29:07.241612Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Data Augmentation\nclass TrainDataset(Dataset):\n    def __init__(self):\n        self.__load_mask_template__()\n        self.__load_train_data__()\n        self.length = self.feature.shape[0]\n        self.label = self.label / 10 - 0.5\n    \n    def __len__(self):\n        return self.length\n    \n    def __getitem__(self, idx):\n        mask = self.mask_template[np.random.choice(len(self.mask_template), 1)[0]]\n        origin_feature = self.feature[idx]\n        label = self.label[idx]\n        feature = np.where(mask == 1, np.nan, origin_feature)\n        imp_mean = SimpleImputer(missing_values=np.nan, strategy='mean')\n        feature = imp_mean.fit_transform(feature.reshape(-1, 1)).reshape(feature.shape)\n        feature = feature / np.max(np.abs(feature))\n        origin_feature = origin_feature / np.max(np.abs(origin_feature))\n        X = np.stack([feature, mask], axis=0)\n        y = np.stack([label, origin_feature[3:-3, 3:-3]], axis=0)\n        return X[:, 3:-3, 3:-3], y\n    \n    def __load_mask_template__(self):\n        self.mask_template = []\n        for file in tqdm(glob.glob('/kaggle/input/gpr-max-deep-learning-challenge-1-gdlc-1/Evaluation_Dataset/*.npy')):\n            data = np.load(file)\n            self.mask_template.append(np.isnan(data).astype(int))\n            \n    def __load_train_data__(self):\n        train = []\n        labels = []\n        train_dir = '/kaggle/input/gpr-max-deep-learning-challenge-1-gdlc-1/Training_Bscan/Training_Bscan'\n        label_dir = '/kaggle/input/gpr-max-deep-learning-challenge-1-gdlc-1/Training_Labels/Training_Labels'\n        for i in tqdm(range(len(os.listdir(train_dir)))):\n            data = np.load(os.path.join(train_dir, f'Bscan_{i}.npy'))\n            train.append(data)\n            data = np.load(os.path.join(label_dir , f'Model_{i}.npy'))\n            labels.append(data)\n        train = np.stack(train, axis=0)\n        labels = np.stack(labels, axis=0)\n        self.feature = train\n        self.label = labels","metadata":{"papermill":{"duration":0.022045,"end_time":"2024-10-29T09:20:44.858451","exception":false,"start_time":"2024-10-29T09:20:44.836406","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-11-15T11:29:07.243723Z","iopub.execute_input":"2024-11-15T11:29:07.244621Z","iopub.status.idle":"2024-11-15T11:29:07.258670Z","shell.execute_reply.started":"2024-11-15T11:29:07.244571Z","shell.execute_reply":"2024-11-15T11:29:07.257603Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# mask_template = []\n# for file in tqdm(glob.glob('/kaggle/input/gpr-max-deep-learning-challenge-1-gdlc-1/Evaluation_Dataset/*.npy')):\n#     data = np.load(file)\n#     mask_template.append(np.isnan(data).astype(int))\n    \n# train = []\n# maskes = []\n# labels = []\n# train_dir = '/kaggle/input/gpr-max-deep-learning-challenge-1-gdlc-1/Training_Bscan/Training_Bscan'\n# label_dir = '/kaggle/input/gpr-max-deep-learning-challenge-1-gdlc-1/Training_Labels/Training_Labels'\n# for i in tqdm(range(len(os.listdir(train_dir)))):\n#     data = np.load(os.path.join(train_dir, f'Bscan_{i}.npy'))\n#     mask = mask_template[np.random.choice(len(mask_template), 1)[0]]\n#     data = np.where(mask == 1, np.nan, data)\n#     train.append(data)\n#     maskes.append(mask)\n#     data = np.load(os.path.join(label_dir , f'Model_{i}.npy'))\n#     labels.append(data)\n    \n# train = np.stack(train, axis=0)\n# maskes = np.stack(maskes, axis=0)\n# labels = np.stack(labels, axis=0)\n# # Load the labels (permittivity ground truth)\n# # np.shape(labels) = (number of samples, 224, 224)\n# # labels = np.load('Labels_Kaggle.npy')\n\n\n# for i, c in enumerate(train):\n#     # Go through my training data i.e. BScans\n#     # Normalise them from -1 to 1\n#     imp_mean = SimpleImputer(missing_values=np.nan, strategy='mean')\n#     c = imp_mean.fit_transform(c.reshape(-1, 1)).reshape(c.shape)\n#     train[i] = c/np.max(np.abs(c))\n#     # Normalise the labels from -0.5 to 0.5\n#     labels[i] = labels[i]/10 - 0.5\n\n#     # plt.imshow(labels[i])\n#     # plt.colorbar()\n#     # plt.show()\n\n# train = np.stack([train, maskes], axis=-1)\n# train = train.transpose((0, 3, 1, 2))\n# # crop\n# train = train[:, :, 3:-3, 3:-3]\n# labels = labels[:, np.newaxis, ...]","metadata":{"papermill":{"duration":0.013852,"end_time":"2024-10-29T09:20:44.877980","exception":false,"start_time":"2024-10-29T09:20:44.864128","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-11-15T11:29:07.261572Z","iopub.execute_input":"2024-11-15T11:29:07.261983Z","iopub.status.idle":"2024-11-15T11:29:07.296821Z","shell.execute_reply.started":"2024-11-15T11:29:07.261936Z","shell.execute_reply":"2024-11-15T11:29:07.295809Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def create_dataloaders(batch_size=32, n_fold=0):\n    # Initialize dataset\n#     dataset = TensorDataset(torch.tensor(train), torch.tensor(labels))\n    dataset = TrainDataset()\n    \n    # Create train/validation split\n    kf = KFold(n_splits=5, shuffle=True, random_state=2024)\n    for i, (train_index, val_index) in enumerate(kf.split(dataset)):\n        if i == n_fold:\n            break\n            \n    train_dataset = Subset(dataset, train_index)\n    val_dataset = Subset(dataset, val_index)\n    print(len(train_dataset))\n    print(len(val_dataset))\n\n    # Create DataLoaders\n    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True)\n    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4, pin_memory=True)\n    \n    return train_loader, val_loader","metadata":{"papermill":{"duration":0.016291,"end_time":"2024-10-29T09:20:44.899951","exception":false,"start_time":"2024-10-29T09:20:44.883660","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-11-15T11:29:07.297920Z","iopub.execute_input":"2024-11-15T11:29:07.298258Z","iopub.status.idle":"2024-11-15T11:29:07.310689Z","shell.execute_reply.started":"2024-11-15T11:29:07.298223Z","shell.execute_reply":"2024-11-15T11:29:07.309750Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_one_epoch(model, train_loader, criterion, optimizer, scheduler, device, reconstruct, deep_supervision):\n    model.train()\n    running_loss = 0.0\n    for images, labels in tqdm(train_loader):\n        if reconstruct == False:\n            labels = labels[:, :1, :, :]\n        images, labels = images.to(device).float(), labels.to(device).float()\n\n        optimizer.zero_grad()\n        \n        outputs = model(images)\n        if deep_supervision:\n            assert outputs.dim() == 5\n            ds_num = outputs.shape[1]\n            if reconstruct:\n                major_loss = criterion(outputs[:, 0, 0, ...], labels[:, 0, ...]) + criterion(outputs[:, 0, 1, ...], labels[:, 1, ...]) * 0.5\n                minor_loss = criterion(outputs[:, 1:, 0, ...], torch.stack([labels[:, 0, ...] for _ in range(ds_num - 1)], dim=1)) + \\\n                            criterion(outputs[:, 1:, 1, ...], torch.stack([labels[:, 1, ...] for _ in range(ds_num - 1)], dim=1)) * 0.5\n            else:\n                major_loss = criterion(outputs[:, 0, ...], labels)\n                minor_loss = criterion(outputs[:, 1:, ...], torch.stack([labels for _ in range(ds_num - 1)], dim=1))\n            loss = major_loss + minor_loss * 0.5\n        else:\n            loss = criterion(outputs, labels)\n\n        # Backward pass and optimization step\n        loss.backward()\n        optimizer.step()\n        scheduler.step()\n\n        running_loss += loss.item() * images.size(0)\n\n    epoch_loss = running_loss / len(train_loader.dataset)\n    return epoch_loss\n\ndef validate(model, val_loader, criterion, device, reconstruct):\n    model.eval()\n    running_loss = 0.0\n    all_labels = []\n    all_outputs = []\n    \n    with torch.no_grad():\n        for images, labels in tqdm(val_loader):\n            if reconstruct == False:\n                labels = labels[:, :1, :, :]\n            images, labels = images.to(device).float(), labels.to(device).float()\n\n            outputs = model(images)\n            loss = criterion(outputs, labels)\n\n            running_loss += loss.item() * images.size(0)\n            all_labels.append(labels.cpu().numpy())\n            all_outputs.append(outputs.cpu().numpy())\n\n    epoch_loss = running_loss / len(val_loader.dataset)\n    \n    all_labels = np.concatenate(all_labels)\n    all_outputs = np.concatenate(all_outputs)\n    \n    return epoch_loss, all_labels, all_outputs","metadata":{"papermill":{"duration":0.020786,"end_time":"2024-10-29T09:20:44.926391","exception":false,"start_time":"2024-10-29T09:20:44.905605","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-11-15T11:29:07.311960Z","iopub.execute_input":"2024-11-15T11:29:07.312369Z","iopub.status.idle":"2024-11-15T11:29:07.329842Z","shell.execute_reply.started":"2024-11-15T11:29:07.312326Z","shell.execute_reply":"2024-11-15T11:29:07.328949Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class EarlyStopping:\n    def __init__(self, patience=7, verbose=False, delta=0):\n        self.patience = patience\n        self.verbose = verbose\n        self.counter = 0\n        self.best_score = None\n        self.early_stop = False\n        self.val_loss_min = np.Inf\n        self.delta = delta\n\n    def __call__(self, val_loss, model, path):\n        score = -val_loss\n        if self.best_score is None:\n            self.best_score = score\n            self.save_checkpoint(val_loss, model, path)\n        elif score < self.best_score + self.delta:\n            self.counter += 1\n            print(f'EarlyStopping counter: {self.counter} out of {self.patience}')\n            if self.counter >= self.patience:\n                self.early_stop = True\n        else:\n            self.best_score = score\n            self.save_checkpoint(val_loss, model, path)\n            self.counter = 0\n\n    def save_checkpoint(self, val_loss, model, path):\n        if self.verbose:\n            print(f'Validation loss decreased ({self.val_loss_min:.6f} --> {val_loss:.6f}).  Saving model ...')\n        torch.save(model.state_dict(), path + '/' + 'checkpoint.pth')\n        self.val_loss_min = val_loss","metadata":{"papermill":{"duration":0.016907,"end_time":"2024-10-29T09:20:44.949002","exception":false,"start_time":"2024-10-29T09:20:44.932095","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-11-15T11:29:07.331082Z","iopub.execute_input":"2024-11-15T11:29:07.331401Z","iopub.status.idle":"2024-11-15T11:29:07.343597Z","shell.execute_reply.started":"2024-11-15T11:29:07.331367Z","shell.execute_reply":"2024-11-15T11:29:07.342757Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_model(model, model_name, num_epochs=10, batch_size=32, lr=1e-4, n_fold=0, device='cuda', patience=3, warmup_epochs=0, reconstruct=False, \n               deep_supervision=False):\n    train_loader, val_loader = create_dataloaders(batch_size=batch_size, n_fold=n_fold)\n    train_sets = len(train_loader)\n\n    model     = model.to(device)\n    criterion = nn.L1Loss()\n    optimizer = optim.Adam(model.parameters(), lr=lr)\n    scheduler = get_cosine_schedule_with_warmup(optimizer, train_sets * warmup_epochs, train_sets * num_epochs)\n    early_stopping = EarlyStopping(patience=patience, verbose=True)\n\n    train_losses = []\n    val_losses   = []\n    \n    path = model_name + str(n_fold)\n    os.makedirs(path, exist_ok=True)\n    \n    for epoch in range(num_epochs):\n        print(f'Epoch {epoch+1}/{num_epochs}')\n        train_loss = train_one_epoch(model, train_loader, criterion, optimizer, scheduler, device, reconstruct, deep_supervision)\n        val_loss, val_labels, val_outputs = validate(model, val_loader, criterion, device, reconstruct)\n\n        train_losses.append(train_loss)\n        val_losses.append(val_loss)\n        \n        # Calculate metrics on validation set\n        mae = mean_absolute_error(val_labels[:, 0, ...].reshape(-1), val_outputs[:, 0, ...].reshape(-1)) \n        print(f'Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}, Val MAE: {mae:.4f}')\n        early_stopping(mae, model, path)\n        if early_stopping.early_stop:\n            print(\"Early stopping\")\n            break\n        print('Updating learning rate to {}'.format(scheduler.get_last_lr()[0]))\n        \n    # Plot Loss\n    plt.plot(train_losses, label='Train Loss')\n    plt.plot(val_losses, label='Validation Loss')\n    plt.legend()\n    plt.savefig('loss_plot.png')\n    plt.show()","metadata":{"papermill":{"duration":0.019057,"end_time":"2024-10-29T09:20:44.973817","exception":false,"start_time":"2024-10-29T09:20:44.954760","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-11-15T11:29:07.344714Z","iopub.execute_input":"2024-11-15T11:29:07.345042Z","iopub.status.idle":"2024-11-15T11:29:07.358719Z","shell.execute_reply.started":"2024-11-15T11:29:07.344975Z","shell.execute_reply":"2024-11-15T11:29:07.357899Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# TODO  \n`filters=[16, 32, 48, 64, 80, 96]`  \n`filters=[32, 64, 96, 128, 160, 192]`  \n`filters=[64, 128, 192, 256, 320, 384]`   \n`deep_supr_num=2`  \n`deep_supervision=False`  \n`res_block=True`  \n`trans_bias=True`  \n`reconstruct=True`  \nimg_size: 224 --> 256","metadata":{"papermill":{"duration":0.005476,"end_time":"2024-10-29T09:20:44.984979","exception":false,"start_time":"2024-10-29T09:20:44.979503","status":"completed"},"tags":[]}},{"cell_type":"code","source":"%%time\nmodel = DynUNet(\n    spatial_dims=2,\n    in_channels=2,\n    out_channels=2,\n    kernel_size=(3, 3, 3, 3, 3, 3),\n    strides=(1, 2, 2, 2, 2, 2),\n    upsample_kernel_size=(2, 2, 2, 2, 2),\n    deep_supervision=True, \n    deep_supr_num=1, \n    dropout=0.2, \n    filters=[64, 96, 128, 192, 256, 384]\n)\nbatch_size = 32\nlr = 3e-3\nreconstruct = True\ntrain_model(model, 'DynUNet', num_epochs=150, batch_size=batch_size, lr=lr, n_fold=0, device='cuda', patience=15, warmup_epochs=15, \n            reconstruct=reconstruct, deep_supervision=True)\ntorch.cuda.empty_cache()\ngc.collect()","metadata":{"papermill":{"duration":4838.024482,"end_time":"2024-10-29T10:41:23.015054","exception":false,"start_time":"2024-10-29T09:20:44.990572","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-11-15T11:29:22.457675Z","iopub.execute_input":"2024-11-15T11:29:22.458096Z","iopub.status.idle":"2024-11-15T13:55:19.034275Z","shell.execute_reply.started":"2024-11-15T11:29:22.458052Z","shell.execute_reply":"2024-11-15T13:55:19.033101Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\ntest_dir = '/kaggle/input/gpr-max-deep-learning-challenge-1-gdlc-1/Evaluation_Dataset'\ntests = []\nfor i in tqdm(range(len(os.listdir(test_dir)))):\n    data = np.load(os.path.join(test_dir, f'Testing_Bscan_{i}.npy'))\n    mask = np.isnan(data).astype(int)\n    imp_mean = SimpleImputer(missing_values=np.nan, strategy='mean')\n    data = imp_mean.fit_transform(data.reshape(-1, 1)).reshape(data.shape)\n    data = data/np.max(np.abs(data))\n    data = np.stack([data, mask], axis=-1)\n    tests.append(data)\ntests = np.stack(tests, axis=0)\ntests = tests.transpose((0, 3, 1, 2))\n# crop\ntests = tests[:, :, 3:-3, 3:-3]","metadata":{"papermill":{"duration":1.288612,"end_time":"2024-10-29T10:41:25.096979","exception":false,"start_time":"2024-10-29T10:41:23.808367","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-11-15T14:04:23.247131Z","iopub.execute_input":"2024-11-15T14:04:23.248102Z","iopub.status.idle":"2024-11-15T14:04:23.758758Z","shell.execute_reply.started":"2024-11-15T14:04:23.248057Z","shell.execute_reply":"2024-11-15T14:04:23.757811Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def predict(model, model_name, batch_size=32, n_fold=0, device='cuda', delete=False, reconstruct=False):\n    model     = model.to(device) # load the model into the GPU\n    model.load_state_dict(torch.load(os.path.join(model_name + str(n_fold), 'checkpoint.pth')))\n    criterion = nn.L1Loss()\n    \n    test_dataset = TensorDataset(torch.tensor(tests), torch.tensor(tests))\n    print(len(test_dataset))\n    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4, pin_memory=True)\n\n    _, _, outputs = validate(model, test_loader, criterion, device, reconstruct)\n    if delete:\n        shutil.rmtree(model_name + str(n_fold))\n    return outputs","metadata":{"papermill":{"duration":0.80527,"end_time":"2024-10-29T10:41:26.742833","exception":false,"start_time":"2024-10-29T10:41:25.937563","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-11-15T14:04:27.540954Z","iopub.execute_input":"2024-11-15T14:04:27.541798Z","iopub.status.idle":"2024-11-15T14:04:27.549379Z","shell.execute_reply.started":"2024-11-15T14:04:27.541759Z","shell.execute_reply":"2024-11-15T14:04:27.547983Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\npreds = predict(model, 'DynUNet', batch_size=batch_size, n_fold=0, device='cuda', delete=False, reconstruct=reconstruct)\npreds = preds[:, 0, ...].reshape(100, -1)\npreds = (preds + 0.5) * 10\nprint(preds.min(), preds.max())\ndf = pd.DataFrame(preds)\ndf[\"Id\"] = range(len(df))\ndf.set_index('Id', inplace=True)\ndf.to_csv('submission.csv')","metadata":{"papermill":{"duration":21.075324,"end_time":"2024-10-29T10:41:48.648005","exception":false,"start_time":"2024-10-29T10:41:27.572681","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-11-15T14:04:32.640297Z","iopub.execute_input":"2024-11-15T14:04:32.641312Z","iopub.status.idle":"2024-11-15T14:04:53.629909Z","shell.execute_reply.started":"2024-11-15T14:04:32.641257Z","shell.execute_reply":"2024-11-15T14:04:53.628714Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"papermill":{"duration":0.788934,"end_time":"2024-10-29T10:41:50.274333","exception":false,"start_time":"2024-10-29T10:41:49.485399","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]}]}