{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# imports","metadata":{}},{"cell_type":"code","source":"import os\nenvironment = os.environ.get(\"KAGGLE_KERNEL_RUN_TYPE\", \"\")","metadata":{"execution":{"iopub.status.busy":"2023-08-10T16:14:01.593457Z","iopub.execute_input":"2023-08-10T16:14:01.593828Z","iopub.status.idle":"2023-08-10T16:14:01.607656Z","shell.execute_reply.started":"2023-08-10T16:14:01.593784Z","shell.execute_reply":"2023-08-10T16:14:01.606675Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import sys\n\nif environment:\n    !pip install -q --no-index --find-links=\"/kaggle/input/segmentation-models-pytorch\" segmentation_models_pytorch \n    sys.path.append(\"/kaggle/input/timm-pretrained-resnest/resnest/\")\n    !pip install -q /kaggle/input/torchsummary/torchsummary-1.5.1-py3-none-any.whl","metadata":{"execution":{"iopub.status.busy":"2023-08-10T16:14:01.84521Z","iopub.execute_input":"2023-08-10T16:14:01.8456Z","iopub.status.idle":"2023-08-10T16:14:34.580144Z","shell.execute_reply.started":"2023-08-10T16:14:01.845571Z","shell.execute_reply":"2023-08-10T16:14:34.574527Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\n\nimport pytorch_lightning as pl\nimport segmentation_models_pytorch as smp\nfrom torch.optim import AdamW\n\nfrom torchmetrics.functional import dice\n\nfrom torch.optim.lr_scheduler import CosineAnnealingLR, ReduceLROnPlateau\nfrom transformers import get_cosine_with_hard_restarts_schedule_with_warmup\n\nimport torchvision.transforms as T\n# import torchvision.transforms.functional as TTF\n\nimport warnings\n\nwarnings.filterwarnings(\"ignore\")\n\nimport os\nimport yaml\nimport numpy as np\nimport pandas as pd\nimport pytorch_lightning as pl\nfrom pprint import pprint\nfrom pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping, TQDMProgressBar\nfrom torch.utils.data import DataLoader\n\n\nimport matplotlib.pyplot as plt\nimport seaborn as sns\n\nfrom torchsummary import summary","metadata":{"execution":{"iopub.status.busy":"2023-08-10T16:14:34.584515Z","iopub.execute_input":"2023-08-10T16:14:34.585214Z","iopub.status.idle":"2023-08-10T16:15:02.862852Z","shell.execute_reply.started":"2023-08-10T16:14:34.58518Z","shell.execute_reply":"2023-08-10T16:15:02.861875Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if environment: \n    !mkdir -p /root/.cache/torch/hub/checkpoints/\n#     !cp /kaggle/input/timm-pretrained-resnest/resnest/gluon_resnest26-50eb607c.pth /root/.cache/torch/hub/checkpoints/gluon_resnest26-50eb607c.pth\n#     !cp /kaggle/input/timm-pretrained-resnest/resnest/resnest50-528c19ca.pth /root/.cache/torch/hub/checkpoints/resnest50-528c19ca.pth","metadata":{"execution":{"iopub.status.busy":"2023-08-10T16:15:02.864299Z","iopub.execute_input":"2023-08-10T16:15:02.864621Z","iopub.status.idle":"2023-08-10T16:15:03.882239Z","shell.execute_reply.started":"2023-08-10T16:15:02.864589Z","shell.execute_reply":"2023-08-10T16:15:03.880912Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Configs","metadata":{}},{"cell_type":"code","source":"%%writefile config.yaml\n\ndata_path: \"/kaggle/input/contrails-images-ash-color\"\noutput_dir: \"models\"\ntrain: True\n\nseed: 420\n\ntrain_bs: 32\nvalid_bs: 128\nworkers: 2\n\nprogress_bar_refresh_rate: 1\n\nearly_stop:\n    monitor: \"val_loss\"\n    mode: \"min\"\n    patience: 999\n    verbose: 1\n\ntrainer:\n    max_epochs: 32\n    min_epochs: 29\n    enable_progress_bar: True\n    precision: \"16-mixed\"\n    devices: 2\n    \n\nmodel:\n    \n    seg_model: \"Unet\"\n    encoder_name: 'timm-resnest50d'\n    loss_smooth: 1.0\n    image_size: 384\n    optimizer_params:\n        lr: 0.0008\n        weight_decay: 0.0001\n    scheduler:\n        name: \"CosineAnnealingLR\"\n        params:\n            CosineAnnealingLR:\n                T_max: 2\n                eta_min: 1.0e-6\n                last_epoch: -1\n            ReduceLROnPlateau:\n                mode: \"min\"\n                factor: 0.31622776602\n                patience: 4\n                verbose: True\n            cosine_with_hard_restarts_schedule_with_warmup:\n                num_warmup_steps: 350\n                num_training_steps: 31500\n                num_cycles: 1","metadata":{"execution":{"iopub.status.busy":"2023-08-10T16:15:03.885222Z","iopub.execute_input":"2023-08-10T16:15:03.885549Z","iopub.status.idle":"2023-08-10T16:15:03.897883Z","shell.execute_reply.started":"2023-08-10T16:15:03.885521Z","shell.execute_reply":"2023-08-10T16:15:03.893847Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with open(\"config.yaml\", \"r\") as file_obj:\n    config = yaml.safe_load(file_obj)","metadata":{"execution":{"iopub.status.busy":"2023-08-10T16:15:03.899329Z","iopub.execute_input":"2023-08-10T16:15:03.900056Z","iopub.status.idle":"2023-08-10T16:15:03.921034Z","shell.execute_reply.started":"2023-08-10T16:15:03.900019Z","shell.execute_reply":"2023-08-10T16:15:03.919806Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Data","metadata":{}},{"cell_type":"markdown","source":"Used [Ash Color Data](https://www.kaggle.com/datasets/shashwatraman/contrails-images-ash-color) by Shashwat Raman. In this dataset he just saved all the train and validation images in its ash color.","metadata":{}},{"cell_type":"code","source":"class ContrailsDataset(torch.utils.data.Dataset):\n    def __init__(self, df, image_size=256, train=True):\n\n        self.df = df\n        self.trn = train\n        self.normalize_image = T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))\n        self.image_size = image_size\n        \n        if image_size != 256:\n            self.resize_image = T.transforms.Resize(image_size)\n            \n        self.random_rotate = T.RandomRotation(degrees=2)\n        self.random_horizontal_flip = T.RandomHorizontalFlip(p=0.5)\n\n    def __getitem__(self, index):\n        row = self.df.iloc[index]\n        con_path = row.path\n        con = np.load(str(con_path))\n\n        img = con[..., :-1]\n        label = con[..., -1]\n\n        label = torch.tensor(label)\n\n        img = torch.tensor(np.reshape(img, (256, 256, 3))).to(torch.float32).permute(2, 0, 1)\n\n        if self.image_size != 256:\n            img = self.resize_image(img)\n\n        img = self.normalize_image(img)\n        \n        if not self.trn:\n            \n            seed = np.random.randint(0, 2**32 - 1)\n            np.random.seed(seed)\n#             print(img.dtype, label.dtype)\n        \n#             dtype = label.dtype\n            label = label.to(torch.float32)\n            label = label.unsqueeze(0)\n        \n            img = self.random_rotate(img)\n            label = self.random_rotate(label)\n            \n            np.random.seed(seed)\n            img = self.random_horizontal_flip(img)\n            label = self.random_horizontal_flip(label)\n            \n            label = label.squeeze(0)\n            \n            label = label.to(torch.float16)\n\n\n        return img.float(), label.float()\n\n    def __len__(self):\n        return len(self.df)","metadata":{"execution":{"iopub.status.busy":"2023-08-10T16:15:03.922867Z","iopub.execute_input":"2023-08-10T16:15:03.923153Z","iopub.status.idle":"2023-08-10T16:15:03.938781Z","shell.execute_reply.started":"2023-08-10T16:15:03.923128Z","shell.execute_reply":"2023-08-10T16:15:03.937583Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\nif config[\"train\"]:\n    contrails = os.path.join(config[\"data_path\"], \"contrails/\")\n    train_path = os.path.join(config[\"data_path\"], \"train_df.csv\")\n    valid_path = os.path.join(config[\"data_path\"], \"valid_df.csv\")\n\n    train_df = pd.read_csv(train_path)\n    valid_df = pd.read_csv(valid_path)\n\n    train_df[\"path\"] = contrails + train_df[\"record_id\"].astype(str) + \".npy\"\n    valid_df[\"path\"] = contrails + valid_df[\"record_id\"].astype(str) + \".npy\"\n\n    dataset_train = ContrailsDataset(train_df, config[\"model\"][\"image_size\"], train=True)\n    dataset_validation = ContrailsDataset(valid_df, config[\"model\"][\"image_size\"], train=False)\n\n    data_loader_train = DataLoader(\n        dataset_train,\n        batch_size=config[\"train_bs\"],\n        shuffle=True,\n        num_workers=config[\"workers\"],\n    )\n    data_loader_validation = DataLoader(\n        dataset_validation,\n        batch_size=config[\"valid_bs\"],\n        shuffle=False,\n        num_workers=config[\"workers\"],\n    )","metadata":{"execution":{"iopub.status.busy":"2023-08-10T16:15:03.940667Z","iopub.execute_input":"2023-08-10T16:15:03.941081Z","iopub.status.idle":"2023-08-10T16:15:03.958638Z","shell.execute_reply.started":"2023-08-10T16:15:03.941048Z","shell.execute_reply":"2023-08-10T16:15:03.957662Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"code","source":"seg_models = {\n    \"Unet\": smp.Unet,\n    \"Unet++\": smp.UnetPlusPlus,\n    \"MAnet\": smp.MAnet,\n    \"Linknet\": smp.Linknet,\n    \"FPN\": smp.FPN,\n    \"PSPNet\": smp.PSPNet,\n    \"PAN\": smp.PAN,\n    \"DeepLabV3\": smp.DeepLabV3,\n    \"DeepLabV3+\": smp.DeepLabV3Plus,\n}\n\n\nclass LightningModule(pl.LightningModule):\n    def __init__(self, config):\n        super().__init__()\n        self.config = config\n        self.model = model = seg_models[config[\"seg_model\"]](\n            encoder_name=config[\"encoder_name\"],\n            encoder_weights=\"imagenet\",\n            in_channels=3,\n            classes=1,\n            activation=None,\n        )\n        \n        self.loss_module = smp.losses.DiceLoss(mode=\"binary\", smooth=config[\"loss_smooth\"])\n        self.val_step_outputs = []\n        self.val_step_labels = []\n\n    def forward(self, batch):\n        imgs = batch\n        preds = self.model(imgs)\n        return preds\n\n    def configure_optimizers(self):\n        optimizer = AdamW(self.parameters(), **self.config[\"optimizer_params\"])\n\n        if self.config[\"scheduler\"][\"name\"] == \"CosineAnnealingLR\":\n            scheduler = CosineAnnealingLR(\n                optimizer,\n                **self.config[\"scheduler\"][\"params\"][\"CosineAnnealingLR\"],\n            )\n            lr_scheduler_dict = {\"scheduler\": scheduler, \"interval\": \"step\"}\n            return {\"optimizer\": optimizer, \"lr_scheduler\": lr_scheduler_dict}\n        elif self.config[\"scheduler\"][\"name\"] == \"ReduceLROnPlateau\":\n            scheduler = ReduceLROnPlateau(\n                optimizer,\n                **self.config[\"scheduler\"][\"params\"][\"ReduceLROnPlateau\"],\n            )\n            lr_scheduler = {\"scheduler\": scheduler, \"monitor\": \"val_loss\"}\n            return {\"optimizer\": optimizer, \"lr_scheduler\": lr_scheduler}\n        elif self.config[\"scheduler\"][\"name\"] == \"cosine_with_hard_restarts_schedule_with_warmup\":\n            scheduler = get_cosine_with_hard_restarts_schedule_with_warmup(\n                optimizer,\n                **self.config[\"scheduler\"][\"params\"][\"cosine_with_hard_restarts_schedule_with_warmup\"],\n            )\n            lr_scheduler_dict = {\"scheduler\": scheduler, \"interval\": \"step\"}\n            return {\"optimizer\": optimizer, \"lr_scheduler\": lr_scheduler_dict}\n        \n\n    def training_step(self, batch, batch_idx):\n        imgs, labels = batch\n        preds = self.model(imgs)\n        if self.config[\"image_size\"] != 256:\n            preds = torch.nn.functional.interpolate(preds, size=256, mode='bilinear')\n        loss = self.loss_module(preds, labels)\n        self.log(\"train_loss\", loss, on_step=True, on_epoch=True, prog_bar=True, batch_size=16)\n\n        for param_group in self.trainer.optimizers[0].param_groups:\n            lr = param_group[\"lr\"]\n        self.log(\"lr\", lr, on_step=True, on_epoch=False, prog_bar=True)\n\n        return loss\n\n    def validation_step(self, batch, batch_idx):\n        imgs, labels = batch\n        preds = self.model(imgs)\n        if self.config[\"image_size\"] != 256:\n            preds = torch.nn.functional.interpolate(preds, size=256, mode='bilinear')\n        loss = self.loss_module(preds, labels)\n        self.log(\"val_loss\", loss, on_step=False, on_epoch=True, prog_bar=True)\n        self.val_step_outputs.append(preds)\n        self.val_step_labels.append(labels)\n\n    def on_validation_epoch_end(self):\n        all_preds = torch.cat(self.val_step_outputs)\n        all_labels = torch.cat(self.val_step_labels)\n        all_preds = torch.sigmoid(all_preds)\n        self.val_step_outputs.clear()\n        self.val_step_labels.clear()\n        val_dice = dice(all_preds, all_labels.long())\n        self.log(\"val_dice\", val_dice, on_step=False, on_epoch=True, prog_bar=True)\n        if self.trainer.global_rank == 0:\n            print(f\"\\nEpoch: {self.current_epoch}\", flush=True)","metadata":{"execution":{"iopub.status.busy":"2023-08-10T16:15:03.960186Z","iopub.execute_input":"2023-08-10T16:15:03.960912Z","iopub.status.idle":"2023-08-10T16:15:03.983685Z","shell.execute_reply.started":"2023-08-10T16:15:03.960878Z","shell.execute_reply":"2023-08-10T16:15:03.982869Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training","metadata":{}},{"cell_type":"code","source":"# Actual training\nif config[\"train\"]:\n\n    checkpoint_callback = ModelCheckpoint(\n        save_weights_only=True,\n        monitor=\"val_dice\",\n        dirpath=config[\"output_dir\"],\n        mode=\"max\",\n        filename=\"model\",\n        save_top_k=1,\n        verbose=1,\n    )\n\n    progress_bar_callback = TQDMProgressBar(\n        refresh_rate=config[\"progress_bar_refresh_rate\"]\n    )\n\n    early_stop_callback = EarlyStopping(**config[\"early_stop\"])\n\n    trainer = pl.Trainer(\n        callbacks=[checkpoint_callback, early_stop_callback, progress_bar_callback],\n        **config[\"trainer\"],\n    )\n\n    config[\"model\"][\"scheduler\"][\"params\"][\"CosineAnnealingLR\"][\"T_max\"] *= len(data_loader_train)/config[\"trainer\"][\"devices\"]\n    model = LightningModule(config[\"model\"])\n    \n\n    trainer.fit(model, data_loader_train, data_loader_validation)","metadata":{"execution":{"iopub.status.busy":"2023-08-10T16:15:03.984842Z","iopub.execute_input":"2023-08-10T16:15:03.98537Z","iopub.status.idle":"2023-08-10T16:15:04.001704Z","shell.execute_reply.started":"2023-08-10T16:15:03.985338Z","shell.execute_reply":"2023-08-10T16:15:04.000651Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(model)","metadata":{"execution":{"iopub.status.busy":"2023-08-10T16:15:08.663298Z","iopub.execute_input":"2023-08-10T16:15:08.663661Z","iopub.status.idle":"2023-08-10T16:15:08.673982Z","shell.execute_reply.started":"2023-08-10T16:15:08.663628Z","shell.execute_reply":"2023-08-10T16:15:08.672887Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Call the trained model","metadata":{}},{"cell_type":"code","source":"class LightningModule(pl.LightningModule):\n\n    def __init__(self):\n        super().__init__()\n        self.model = smp.Unet(encoder_name=config[\"model\"][\"encoder_name\"],\n                              encoder_weights=None,\n                              decoder_use_batchnorm=True,\n                              in_channels=3,\n                              classes=1,\n                              activation=None,\n                              )\n\n    def forward(self, batch):\n        return self.model(batch)\n\n    \nif config[\"train\"]:\n    model = LightningModule().load_from_checkpoint(\"/kaggle/working/models/model.ckpt\")\nelse:\n    model = LightningModule().load_from_checkpoint(\"/kaggle/input/models/models/model.ckpt\")\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nmodel.to(device)\nmodel.eval()\nmodel.zero_grad()","metadata":{"execution":{"iopub.status.busy":"2023-08-10T16:15:23.497548Z","iopub.execute_input":"2023-08-10T16:15:23.497958Z","iopub.status.idle":"2023-08-10T16:15:36.431327Z","shell.execute_reply.started":"2023-08-10T16:15:23.497925Z","shell.execute_reply":"2023-08-10T16:15:36.430204Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Validation and Test data load","metadata":{}},{"cell_type":"code","source":"class ContrailsDataset(torch.utils.data.Dataset):\n    def __init__(self, df, image_size=256, train=True, transforms=None):\n        self.df = df  # Initialize the instance variable df to store the DataFrame.\n        self.trn = train  # Initialize the instance variable trn to indicate if it is a training dataset.\n        self.transforms = transforms  # Initialize the instance variable transforms to store the transforms.\n        self.image_size = image_size\n        if image_size != 256:\n            self.resize_image = T.transforms.Resize(image_size)\n        \n        self.normalize_image = T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))\n        \n\n    def read_record(self, directory):\n\n        record_data = {}  # Create a dictionary to store the record data.\n        for x in [\n            \"band_11\",\n            \"band_14\",\n            \"band_15\"\n        ]:\n            record_data[x] = np.load(os.path.join(directory, x + \".npy\"))  # Load data for each band and store it in the dictionary.\n\n        if self.trn:\n            record_data[\"mask\"] = np.load(os.path.join(directory, \"human_pixel_masks.npy\"))\n\n        return record_data\n\n    def normalize_range(self, data, bounds):\n        \"\"\"Normalize data to the range [0, 1].\"\"\"\n        return (data - bounds[0]) / (bounds[1] - bounds[0])\n\n    def get_false_color(self, record_data):\n\n        _T11_BOUNDS = (243, 303)\n        _CLOUD_TOP_TDIFF_BOUNDS = (-4, 5)\n        _TDIFF_BOUNDS = (-4, 2)\n\n        N_TIMES_BEFORE = 4\n        r = self.normalize_range(record_data[\"band_15\"] - record_data[\"band_14\"], _TDIFF_BOUNDS)\n        g = self.normalize_range(record_data[\"band_14\"] - record_data[\"band_11\"], _CLOUD_TOP_TDIFF_BOUNDS)\n        b = self.normalize_range(record_data[\"band_14\"], _T11_BOUNDS)\n        false_color = np.clip(np.stack([r, g, b], axis=2), 0, 1)\n        img = false_color[..., N_TIMES_BEFORE]\n\n        if self.trn:\n            mask_img = record_data[\"mask\"]\n\n            return img, mask_img\n        \n        return img\n\n    def __getitem__(self, index):\n        row = self.df.iloc[index]\n        con_path = row.path\n        data = self.read_record(con_path)  # dictionary with keys: band_11, band_14, band_15 and values: numpy arrays (height, width, channels)\n\n        if self.trn:\n            img, mask_img = self.get_false_color(data)\n\n            img = torch.tensor(img).float()\n            mask_img = torch.tensor(mask_img).float()\n\n            img = img.permute(2, 0, 1)\n            mask_img = mask_img.permute(2, 0, 1)\n            \n            if self.image_size != 256:\n                img = self.resize_image(img)\n            \n            img = self.normalize_image(img)\n\n            return img, mask_img\n        \n        img = self.get_false_color(data)\n        \n        img = torch.tensor(img).float()\n\n        img = img.permute(2, 0, 1)\n        \n        if self.image_size != 256:\n            img = self.resize_image(img)\n        \n        img = self.normalize_image(img)\n        \n        return img\n    \n\n    def __len__(self):\n        return len(self.df)\n","metadata":{"execution":{"iopub.status.busy":"2023-08-10T16:15:39.533794Z","iopub.execute_input":"2023-08-10T16:15:39.534221Z","iopub.status.idle":"2023-08-10T16:15:39.557774Z","shell.execute_reply.started":"2023-08-10T16:15:39.53419Z","shell.execute_reply":"2023-08-10T16:15:39.556481Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_paths(data_type):\n    \n    base_dir = \"/kaggle/input/google-research-identify-contrails-reduce-global-warming\"\n\n    ids_list = os.listdir(os.path.join(base_dir, data_type))\n\n    df = pd.DataFrame(ids_list, columns=['record_id'])\n\n    df['path'] = os.path.join(base_dir, data_type ) +\"/\"+ df['record_id'].astype(str)\n\n    return df\n\nval_df = get_paths('validation')\n\nval_ds = ContrailsDataset(\n        val_df,\n        config[\"model\"][\"image_size\"],\n        train = True\n    )\n \nval_dl = DataLoader(val_ds, batch_size=8, shuffle=True, num_workers = 2)","metadata":{"execution":{"iopub.status.busy":"2023-08-10T16:15:49.525704Z","iopub.execute_input":"2023-08-10T16:15:49.526421Z","iopub.status.idle":"2023-08-10T16:15:49.695931Z","shell.execute_reply.started":"2023-08-10T16:15:49.526388Z","shell.execute_reply":"2023-08-10T16:15:49.694843Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Optimum Threshold","metadata":{}},{"cell_type":"code","source":"class DiceThresholdTester:\n    \n    def __init__(self, model: nn.Module, data_loader: torch.utils.data.DataLoader):\n        self.model = model\n        self.data_loader = data_loader\n        self.cumulative_mask_pred = []\n        self.cumulative_mask_true = []\n        \n    def precalculate_prediction(self) -> None:\n        sigmoid = nn.Sigmoid()\n        \n        for images, mask_true in self.data_loader:\n            if torch.cuda.is_available():\n                images = images.cuda()\n\n            mask_pred = sigmoid(model.forward(images))\n            \n            if config[\"model\"][\"image_size\"] != 256:\n                mask_pred = torch.nn.functional.interpolate(mask_pred, size=256, mode='bilinear')\n\n            self.cumulative_mask_pred.append(mask_pred.cpu().detach().numpy())\n            self.cumulative_mask_true.append(mask_true.cpu().detach().numpy())\n\n            \n        self.cumulative_mask_pred = np.concatenate(self.cumulative_mask_pred, axis=0)\n        self.cumulative_mask_true = np.concatenate(self.cumulative_mask_true, axis=0)\n\n        self.cumulative_mask_pred = torch.flatten(torch.from_numpy(self.cumulative_mask_pred))\n        self.cumulative_mask_true = torch.flatten(torch.from_numpy(self.cumulative_mask_true))\n    \n    def test_threshold(self, threshold: float) -> float:\n        after_threshold = np.zeros(self.cumulative_mask_pred.shape)\n        after_threshold[self.cumulative_mask_pred[:] > threshold] = 1\n        after_threshold[self.cumulative_mask_pred[:] < threshold] = 0\n        after_threshold = torch.flatten(torch.from_numpy(after_threshold))\n\n        return dice(after_threshold, self.cumulative_mask_true.long())","metadata":{"execution":{"iopub.status.busy":"2023-08-10T16:15:57.973783Z","iopub.execute_input":"2023-08-10T16:15:57.97422Z","iopub.status.idle":"2023-08-10T16:15:57.986355Z","shell.execute_reply.started":"2023-08-10T16:15:57.974187Z","shell.execute_reply":"2023-08-10T16:15:57.985064Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if config[\"train\"]:\n    dice_threshold_tester = DiceThresholdTester(model, val_dl)\n    dice_threshold_tester.precalculate_prediction()","metadata":{"execution":{"iopub.status.busy":"2023-08-10T16:16:14.676447Z","iopub.execute_input":"2023-08-10T16:16:14.67683Z","iopub.status.idle":"2023-08-10T16:18:24.916556Z","shell.execute_reply.started":"2023-08-10T16:16:14.676786Z","shell.execute_reply":"2023-08-10T16:18:24.914974Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if config[\"train\"]:\n    thresholds_to_test = [round(x * 0.01, 2) for x in range(101)]\n\n    optim_threshold = 0.98\n    best_dice_score = -1\n\n    thresholds = []\n    dice_scores = []\n\n    for t in thresholds_to_test:\n        dice_score = dice_threshold_tester.test_threshold(t)\n        if dice_score > best_dice_score:\n            best_dice_score = dice_score\n            optim_threshold = t\n\n        thresholds.append(t)\n        dice_scores.append(dice_score)\n\n    print(f'Best Threshold: {optim_threshold} with dice: {best_dice_score}')\n    df_threshold_data = pd.DataFrame({'Threshold': thresholds, 'Dice Score': dice_scores})\nelse:\n    optim_threshold = 0.25","metadata":{"execution":{"iopub.status.busy":"2023-08-10T16:18:49.476339Z","iopub.execute_input":"2023-08-10T16:18:49.47678Z","iopub.status.idle":"2023-08-10T16:31:21.68412Z","shell.execute_reply.started":"2023-08-10T16:18:49.476744Z","shell.execute_reply":"2023-08-10T16:31:21.682891Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Vizualization","metadata":{}},{"cell_type":"code","source":"batches_to_show = 4\n\nfor i, data in enumerate(val_dl):\n    images, mask = data\n    \n    # Predict mask for this instance\n    if torch.cuda.is_available():\n        images = images.cuda()\n    \n        \n    predicated_mask = torch.sigmoid(model.forward(images[:, :, :, :]))\n    \n    if config[\"model\"][\"image_size\"] != 256:\n        predicated_mask = torch.nn.functional.interpolate(predicated_mask, size=256, mode='bilinear').cpu().detach().numpy()\n    else:\n        predicated_mask = predicated_mask.cpu().detach().numpy()\n\n    \n    # Apply threshold\n    predicated_mask_with_threshold = np.zeros((images.shape[0], 256, 256))\n    predicated_mask_with_threshold[predicated_mask[:, 0, :, :] < optim_threshold] = 0\n    predicated_mask_with_threshold[predicated_mask[:, 0, :, :] > optim_threshold] = 1\n    \n    images = images.cpu()\n        \n    for img_num in range(0, images.shape[0]):\n        fig, axes = plt.subplots(nrows=1, ncols=4, figsize=(20,10))\n        axes = axes.flatten()\n        \n        # Show groud trought \n        axes[0].imshow(mask[img_num, 0, :, :])\n        axes[0].axis('off')\n        axes[0].set_title('Ground Truth')\n        \n        # Show ash color scheme input image\n        axes[1].imshow(images[img_num, :, :, :].permute(1, 2, 0))\n        axes[1].axis('off')\n        axes[1].set_title('Ash color scheeme input - Frame 4')\n\n        # Show predicted mask\n        axes[2].imshow(predicated_mask[img_num, 0, :, :], vmin=0, vmax=1)\n        axes[2].axis('off')\n        axes[2].set_title('Predicted probability mask')\n\n        # Show predicted mask after threshold\n        axes[3].imshow(predicated_mask_with_threshold[img_num, :, :])\n        axes[3].axis('off')\n        axes[3].set_title('Predicted mask with threshold')\n        plt.show()\n    \n    if i + 1 >= batches_to_show:\n        break","metadata":{"execution":{"iopub.status.busy":"2023-08-10T16:33:16.660415Z","iopub.execute_input":"2023-08-10T16:33:16.660899Z","iopub.status.idle":"2023-08-10T16:33:38.376495Z","shell.execute_reply.started":"2023-08-10T16:33:16.660858Z","shell.execute_reply":"2023-08-10T16:33:38.375181Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Submission","metadata":{}},{"cell_type":"code","source":"torch.cuda.empty_cache()","metadata":{"execution":{"iopub.status.busy":"2023-08-10T16:33:41.847062Z","iopub.execute_input":"2023-08-10T16:33:41.847495Z","iopub.status.idle":"2023-08-10T16:33:42.01603Z","shell.execute_reply.started":"2023-08-10T16:33:41.847465Z","shell.execute_reply":"2023-08-10T16:33:42.014996Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_df = get_paths('test')\n\n# cast record_id to int\ntest_df[\"record_id\"] = test_df.record_id.astype(int)\n\ntest_ds = ContrailsDataset(\n        test_df,\n        config[\"model\"][\"image_size\"],\n        train = False\n    )\n\ntest_batch_size = 1\n\ntest_dl = DataLoader(test_ds, batch_size=test_batch_size, num_workers = config[\"workers\"])\n\ndel test_ds","metadata":{"execution":{"iopub.status.busy":"2023-08-10T16:33:43.173584Z","iopub.execute_input":"2023-08-10T16:33:43.17399Z","iopub.status.idle":"2023-08-10T16:33:43.197103Z","shell.execute_reply.started":"2023-08-10T16:33:43.173958Z","shell.execute_reply":"2023-08-10T16:33:43.196061Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def rle_encode(x, fg_val=1):\n    \"\"\"\n    Args:\n        x:  numpy array of shape (height, width), 1 - mask, 0 - background\n    Returns: run length encoding as list\n    \"\"\"\n\n    dots = np.where(\n        x.T.flatten() == fg_val)[0]  # .T sets Fortran order down-then-right\n    run_lengths = []\n    prev = -2\n    for b in dots:\n        if b > prev + 1:\n            run_lengths.extend((b + 1, 0))\n        run_lengths[-1] += 1\n        prev = b\n    return run_lengths\n\ndef list_to_string(x):\n    \"\"\"\n    Converts list to a string representation\n    Empty list returns '-'\n    \"\"\"\n    if x: # non-empty list\n        s = str(x).replace(\"[\", \"\").replace(\"]\", \"\").replace(\",\", \"\")\n    else:\n        s = '-'\n    return s","metadata":{"execution":{"iopub.status.busy":"2023-08-10T16:33:45.652542Z","iopub.execute_input":"2023-08-10T16:33:45.65295Z","iopub.status.idle":"2023-08-10T16:33:45.66245Z","shell.execute_reply.started":"2023-08-10T16:33:45.652917Z","shell.execute_reply":"2023-08-10T16:33:45.661078Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission = pd.read_csv(os.path.join(\"/kaggle/input/google-research-identify-contrails-reduce-global-warming\", \"sample_submission.csv\"), index_col='record_id')","metadata":{"execution":{"iopub.status.busy":"2023-08-10T16:33:46.543687Z","iopub.execute_input":"2023-08-10T16:33:46.54463Z","iopub.status.idle":"2023-08-10T16:33:46.602102Z","shell.execute_reply.started":"2023-08-10T16:33:46.544593Z","shell.execute_reply":"2023-08-10T16:33:46.601049Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for i, images in enumerate(test_dl):\n    \n    \n    image_id = torch.tensor(test_df.iloc[i]['record_id'])\n    \n    # Predict mask for this instance\n    images = images.to(device)\n    with torch.no_grad():\n        predicted_mask = model.forward(images[:, :, :, :])\n        \n    if config[\"model\"][\"image_size\"] != 256:\n        predicted_mask = torch.nn.functional.interpolate(predicted_mask, size=256, mode='bilinear')\n    predicated_mask = predicted_mask.cpu().detach().numpy()\n    \n    # Apply threshold\n    predicated_mask_with_threshold = np.zeros((images.shape[0], 256, 256))\n    predicated_mask_with_threshold[predicated_mask[:, 0, :, :] < optim_threshold] = 0\n    predicated_mask_with_threshold[predicated_mask[:, 0, :, :] > optim_threshold] = 1\n    \n    current_mask = predicated_mask_with_threshold[:, :, :]\n    current_image_id = image_id.item()\n    submission.loc[int(current_image_id), 'encoded_pixels'] = list_to_string(rle_encode(current_mask))","metadata":{"execution":{"iopub.status.busy":"2023-08-10T16:33:47.513208Z","iopub.execute_input":"2023-08-10T16:33:47.513605Z","iopub.status.idle":"2023-08-10T16:33:48.367256Z","shell.execute_reply.started":"2023-08-10T16:33:47.513574Z","shell.execute_reply":"2023-08-10T16:33:48.365807Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission.head()","metadata":{"execution":{"iopub.status.busy":"2023-08-10T16:33:48.373714Z","iopub.execute_input":"2023-08-10T16:33:48.376144Z","iopub.status.idle":"2023-08-10T16:33:48.413439Z","shell.execute_reply.started":"2023-08-10T16:33:48.376106Z","shell.execute_reply":"2023-08-10T16:33:48.412366Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission.to_csv('submission.csv')","metadata":{"execution":{"iopub.status.busy":"2023-08-10T16:33:49.458563Z","iopub.execute_input":"2023-08-10T16:33:49.458961Z","iopub.status.idle":"2023-08-10T16:33:49.470249Z","shell.execute_reply.started":"2023-08-10T16:33:49.458931Z","shell.execute_reply":"2023-08-10T16:33:49.469204Z"},"trusted":true},"execution_count":null,"outputs":[]}]}