{"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 gc\nimport os\nimport sys\nimport cv2\nimport glob\nimport json\nimport shutil\nimport random\nimport pydicom\nimport numpy as np\nimport pandas as pd\nimport seaborn as sns\nfrom PIL import Image\nfrom tqdm.notebook import tqdm\nimport matplotlib.pyplot as plt\n\nfrom sklearn.metrics import f1_score\nfrom sklearn.preprocessing import LabelEncoder\nfrom sklearn.model_selection import StratifiedKFold, StratifiedGroupKFold\n\nsys.path.append('../input/timm-pytorch-image-models/pytorch-image-models-master/')\nimport timm\nfrom timm.data import create_transform\nfrom timm import create_model, list_models\n\nimport torch\nfrom torch import nn\nimport torch.nn.functional as F\nfrom torch.cuda.amp import GradScaler, autocast\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.utils.data.sampler import Sampler \nfrom torch.nn.functional import binary_cross_entropy_with_logits, cross_entropy\n\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\nfrom transformers import Mask2FormerConfig, Mask2FormerImageProcessor, Mask2FormerForUniversalSegmentation\nfrom transformers import OneFormerConfig, OneFormerImageProcessor, OneFormerForUniversalSegmentation","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-05-24T08:19:36.171349Z","iopub.execute_input":"2023-05-24T08:19:36.171825Z","iopub.status.idle":"2023-05-24T08:19:36.198142Z","shell.execute_reply.started":"2023-05-24T08:19:36.171787Z","shell.execute_reply":"2023-05-24T08:19:36.197279Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## CFG","metadata":{}},{"cell_type":"code","source":"class CFG:\n    \"\"\"\n    Parameters used for training\n    \"\"\"\n    seed = 42\n    \n    img_size = (512, 512)\n    batch_size = 1\n    epochs = 10\n    use_fp16 = True\n    n_folds = 5\n    train_folds = [0, 1, 2, 3, 4]\n    \n    weight_decay = 0.024\n    one_cycle_max_lr = 4e-4  # 8e-4    \n    tta = False\n\n# os.environ['CUDA_VISIBLE_DEVICES'] = \"0,1\"\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")","metadata":{"execution":{"iopub.status.busy":"2023-05-24T07:48:54.946312Z","iopub.execute_input":"2023-05-24T07:48:54.947072Z","iopub.status.idle":"2023-05-24T07:48:54.975954Z","shell.execute_reply.started":"2023-05-24T07:48:54.947032Z","shell.execute_reply":"2023-05-24T07:48:54.974965Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def set_seed(seed = 42):\n    '''Sets the seed of the entire notebook so results are the same every time we run.\n    This is for REPRODUCIBILITY.'''\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    # When running on the CuDNN backend, two further options must be set\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n    # Set a fixed value for the hash seed\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    print('> SEEDING DONE')\n    \nset_seed(CFG.seed)","metadata":{"execution":{"iopub.status.busy":"2023-05-24T07:48:54.977394Z","iopub.execute_input":"2023-05-24T07:48:54.977991Z","iopub.status.idle":"2023-05-24T07:48:54.997203Z","shell.execute_reply.started":"2023-05-24T07:48:54.977956Z","shell.execute_reply":"2023-05-24T07:48:54.996219Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Data Prep","metadata":{}},{"cell_type":"code","source":"COMP_PATH = \"/kaggle/input/hubmap-hacking-the-human-vasculature\"\ntile_meta = pd.read_csv(f\"{COMP_PATH}/tile_meta.csv\")\nwsi_meta = pd.read_csv(f\"{COMP_PATH}/wsi_meta.csv\")","metadata":{"execution":{"iopub.status.busy":"2023-05-24T07:48:54.999534Z","iopub.execute_input":"2023-05-24T07:48:55.001667Z","iopub.status.idle":"2023-05-24T07:48:55.032808Z","shell.execute_reply.started":"2023-05-24T07:48:55.001642Z","shell.execute_reply":"2023-05-24T07:48:55.032006Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with open(f\"{COMP_PATH}/polygons.jsonl\", \"r\") as json_file:\n    json_list = list(json_file)\n    \ntiles_data = {}\nfor json_str in tqdm(json_list, total=len(json_list)):\n    json_data = json.loads(json_str)\n    tiles_data[json_data['id']] = json_data['annotations']","metadata":{"execution":{"iopub.status.busy":"2023-05-24T07:48:55.034076Z","iopub.execute_input":"2023-05-24T07:48:55.034528Z","iopub.status.idle":"2023-05-24T07:48:59.885826Z","shell.execute_reply.started":"2023-05-24T07:48:55.034496Z","shell.execute_reply":"2023-05-24T07:48:59.884799Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_no_of_blood_vessel(img_id):\n    mask = np.zeros(CFG.img_size, dtype=np.float32)\n    cnt = 0\n    for annot in tiles_data[img_id]:\n        if annot['type'] != \"blood_vessel\": continue\n        for cord in annot['coordinates']:\n            row, col = np.array([i[1] for i in cord]), np.asarray([i[0] for i in cord])\n            mask[row, col] = 1\n            cnt += 1\n    return cnt","metadata":{"execution":{"iopub.status.busy":"2023-05-24T07:48:59.887391Z","iopub.execute_input":"2023-05-24T07:48:59.887783Z","iopub.status.idle":"2023-05-24T07:48:59.896202Z","shell.execute_reply.started":"2023-05-24T07:48:59.887738Z","shell.execute_reply":"2023-05-24T07:48:59.894979Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## CV","metadata":{}},{"cell_type":"code","source":"df = tile_meta.query('dataset < 3').reset_index(drop=True)\ndf['no_of_bv'] = df['id'].apply(get_no_of_blood_vessel)\ndf = df.query('no_of_bv > 0').reset_index(drop=True)\n\nskf = StratifiedKFold(n_splits=CFG.n_folds, random_state=CFG.seed, shuffle=True)\ndf['fold'] = -1\n\nfor fold, (train_idx, test_idx) in enumerate(skf.split(df, df['source_wsi'])):\n    df.loc[test_idx, 'fold'] = fold\n\ndf.groupby('fold')['source_wsi'].value_counts()","metadata":{"execution":{"iopub.status.busy":"2023-05-24T07:48:59.897671Z","iopub.execute_input":"2023-05-24T07:48:59.898279Z","iopub.status.idle":"2023-05-24T07:49:00.972387Z","shell.execute_reply.started":"2023-05-24T07:48:59.898247Z","shell.execute_reply":"2023-05-24T07:49:00.971483Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Dataset","metadata":{}},{"cell_type":"code","source":"class HubmapDataset(Dataset):\n    def __init__(self, df, config, transforms=None):\n        self.df = df\n        self.img_ids = self.df.id.values\n        self.transforms = transforms\n        self.processor = Mask2FormerImageProcessor(config)\n    \n    def __len__(self):\n        return len(self.img_ids)\n    \n    def __getitem__(self, index):\n        img_id = self.img_ids[index]\n        img_path = f\"{COMP_PATH}/train/{img_id}.tif\"\n        \n        try:\n            image = np.asarray(Image.open(img_path).convert('RGB'))\n        except Exception as ex:\n            print(img_path, ex)\n            return None\n        \n        mask = np.zeros((CFG.img_size[0], CFG.img_size[1]), dtype=np.float32)\n        cnt = 0\n        for annot in tiles_data[img_id]:\n            if annot['type'] != \"blood_vessel\": continue\n            for cord in annot['coordinates']:\n                row, col = np.array([i[1] for i in cord]), np.asarray([i[0] for i in cord])\n                mask[row, col] = 1\n                cnt += 1\n        \n        if self.transforms is not None:\n            transformed = self.transforms(image=image, mask=mask)\n            image, mask = transformed['image'], torch.unsqueeze(transformed['mask'], 0)\n        \n        try:\n            inputs = self.processor.encode_inputs(pixel_values_list=[image], \n                                                  task_inputs=['instance'],\n                                             segmentation_maps=mask, \n                                             ignore_index=0,\n                                             return_tensors='pt')\n        except Exception as err:\n            print(image.shape, mask.shape)\n            print(img_id, cnt)\n            plt.imshow(np.asarray(Image.open(img_path).convert('RGB')))\n            plt.imshow(mask.permute(1, 2, 0))\n            print(err)\n            plt.show()\n            return {}\n\n        inputs['pixel_values'] = inputs['pixel_values'][0]\n        return inputs","metadata":{"execution":{"iopub.status.busy":"2023-05-24T08:29:12.760500Z","iopub.execute_input":"2023-05-24T08:29:12.760871Z","iopub.status.idle":"2023-05-24T08:29:12.775117Z","shell.execute_reply.started":"2023-05-24T08:29:12.760843Z","shell.execute_reply":"2023-05-24T08:29:12.773712Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Augmentation","metadata":{}},{"cell_type":"code","source":"def transformer(stage):\n    if stage == \"train\":\n        return A.Compose([\n            A.HorizontalFlip(p=0.5),\n            A.Rotate(limit=5),\n#             A.augmentations.crops.RandomResizedCrop(height=CFG.img_size[0], width=CFG.img_size[1], scale=(0.8, 1), ratio=(0.45, 0.55)),\n            A.Normalize(),\n            A.pytorch.transforms.ToTensorV2()\n        ])\n    else:\n        return A.Compose([\n                A.Resize(CFG.img_size[0], CFG.img_size[1]),\n                A.Normalize(),\n                A.pytorch.transforms.ToTensorV2()\n            ])","metadata":{"execution":{"iopub.status.busy":"2023-05-24T07:49:00.989136Z","iopub.execute_input":"2023-05-24T07:49:00.989757Z","iopub.status.idle":"2023-05-24T07:49:00.999782Z","shell.execute_reply.started":"2023-05-24T07:49:00.989725Z","shell.execute_reply":"2023-05-24T07:49:00.998846Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def plot_df(df):\n    fig,ax = plt.subplots(1,1,figsize=(15,5))\n    ax.plot(df['train_loss'])\n    ax.plot(df['valid_loss'])\n    ax.legend()\n    ax.set_title('Loss')","metadata":{"execution":{"iopub.status.busy":"2023-05-24T07:49:01.004411Z","iopub.execute_input":"2023-05-24T07:49:01.004681Z","iopub.status.idle":"2023-05-24T07:49:01.013421Z","shell.execute_reply.started":"2023-05-24T07:49:01.004657Z","shell.execute_reply":"2023-05-24T07:49:01.012273Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model","metadata":{}},{"cell_type":"code","source":"class HubmapModel(nn.Module):\n    def __init__(self, config):\n        super(HubmapModel, self).__init__()\n        self.model = Mask2FormerForUniversalSegmentation(config)\n    \n    def forward(self, inputs):\n        outputs = self.model(**inputs)\n        return outputs","metadata":{"execution":{"iopub.status.busy":"2023-05-24T08:20:35.956823Z","iopub.execute_input":"2023-05-24T08:20:35.957190Z","iopub.status.idle":"2023-05-24T08:20:35.962954Z","shell.execute_reply.started":"2023-05-24T08:20:35.957161Z","shell.execute_reply":"2023-05-24T08:20:35.961749Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"config = Mask2FormerConfig(feature_size=CFG.img_size[0], mask_feature_size=CFG.img_size[0])\nconfig.save_pretrained('./config')","metadata":{"execution":{"iopub.status.busy":"2023-05-24T08:30:49.625282Z","iopub.execute_input":"2023-05-24T08:30:49.625682Z","iopub.status.idle":"2023-05-24T08:30:49.634085Z","shell.execute_reply.started":"2023-05-24T08:30:49.625650Z","shell.execute_reply":"2023-05-24T08:30:49.633133Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def add_weight_decay(model, weight_decay=1e-5, skip_list=()):\n    decay = []\n    no_decay = []\n    for name, param in model.named_parameters():\n        if not param.requires_grad:\n            continue\n        if len(param.shape) == 1 or np.any([v in name.lower() for v in skip_list]):\n            no_decay.append(param)\n        else:\n            decay.append(param)\n    return [\n        {'params': no_decay, 'weight_decay': 0.},\n        {'params': decay, 'weight_decay': weight_decay}]","metadata":{"execution":{"iopub.status.busy":"2023-05-24T08:21:11.250013Z","iopub.execute_input":"2023-05-24T08:21:11.250372Z","iopub.status.idle":"2023-05-24T08:21:11.256486Z","shell.execute_reply.started":"2023-05-24T08:21:11.250343Z","shell.execute_reply":"2023-05-24T08:21:11.255626Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_optimizer_and_scheduler(model, dataloader, optim=\"adam\"):\n    if optim == \"adamw\":\n         optimizer = torch.optim.AdamW(\n             add_weight_decay(model,\n                              weight_decay=CFG.weight_decay,\n                              skip_list=['bias']),\n             lr=CFG.one_cycle_max_lr,\n             betas=(0.9, 0.999),\n             weight_decay=CFG.weight_decay)\n    else:\n        optimizer = torch.optim.Adam(model.parameters())\n    \n    scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer=optimizer, pct_start=0.1,\n                                              max_lr=CFG.one_cycle_max_lr, epochs=CFG.epochs, steps_per_epoch=len(dataloader))\n    return optimizer, scheduler","metadata":{"execution":{"iopub.status.busy":"2023-05-24T08:21:12.504267Z","iopub.execute_input":"2023-05-24T08:21:12.504632Z","iopub.status.idle":"2023-05-24T08:21:12.513198Z","shell.execute_reply.started":"2023-05-24T08:21:12.504600Z","shell.execute_reply":"2023-05-24T08:21:12.511633Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Train","metadata":{}},{"cell_type":"code","source":"def train_one_epoch(dataloader, model, scheduler, optimizer, scaler, epoch):\n    model.train()\n\n    total_loss = 0\n    pbar = tqdm(dataloader, desc=f\"Train: Epoch {epoch + 1}\", total=len(dataloader), mininterval=5)\n\n    for inputs in pbar:\n        optimizer.zero_grad()\n        \n        # Using mixed precision training\n        with autocast():\n            inps = { k: (v.to(device) if type(v) == torch.Tensor else torch.stack(v).to(device)) for k, v in inputs.items() }\n            outputs = model(inps)\n            loss = outputs.loss\n            \n            if np.isinf(loss.item()) or np.isnan(loss.item()):\n                print(f'Bad loss, skipping the batch')\n                del loss, outputs\n                gc_collect()\n                continue\n\n        # scaler is needed to prevent \"gradient underflow\"\n        scaler.scale(loss).backward()\n        scaler.step(optimizer)\n        scale = scaler.get_scale()\n        scaler.update()\n        skip_lr_scheduler = scale > scaler.get_scale()\n        if scheduler is not None and not skip_lr_scheduler:\n            scheduler.step()\n        \n        lr = scheduler.get_last_lr()[0] if scheduler else CFG.one_cycle_max_lr\n        loss = loss.item()\n        \n        pbar.set_postfix({\"loss\": loss, \"lr\": lr})\n        total_loss += loss\n    \n    total_loss /= len(dataloader)\n    gc.collect()\n    torch.cuda.empty_cache()\n    return total_loss","metadata":{"execution":{"iopub.status.busy":"2023-05-24T08:21:43.569595Z","iopub.execute_input":"2023-05-24T08:21:43.569945Z","iopub.status.idle":"2023-05-24T08:21:43.582088Z","shell.execute_reply.started":"2023-05-24T08:21:43.569918Z","shell.execute_reply":"2023-05-24T08:21:43.581196Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def valid_one_epoch(dataloader, model, epoch):\n    model.eval()\n    with torch.no_grad():\n        total_loss = 0\n        pbar = tqdm(dataloader, desc=f'Eval: {epoch + 1}', total=len(dataloader), mininterval=5)\n\n        for inputs in pbar:\n            with autocast(enabled=False):\n                inputs = { k: (v.to(device) if type(v) == torch.Tensor else [v[0].to(device)]) for k, v in inputs.items() }\n                outputs = model(inputs)\n\n                loss = outputs.item()\n                pbar.set_postfix({\"loss\": loss})\n                total_loss += loss\n\n    total_loss /= len(dataloader)\n    gc.collect()\n    torch.cuda.empty_cache()\n    return total_loss","metadata":{"execution":{"iopub.status.busy":"2023-05-24T08:21:47.603511Z","iopub.execute_input":"2023-05-24T08:21:47.604342Z","iopub.status.idle":"2023-05-24T08:21:47.613197Z","shell.execute_reply.started":"2023-05-24T08:21:47.604306Z","shell.execute_reply":"2023-05-24T08:21:47.612222Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_fnc(train_dataloader, valid_dataloader, model, fold, optimizer, scheduler):\n    train_losses = []\n    valid_losses = []\n    \n    scaler = GradScaler()\n    best_loss = 999\n    best_score = -1\n    for epoch in range(CFG.epochs):\n        train_loss = train_one_epoch(train_dataloader, model, scheduler, optimizer, scaler, epoch)\n        valid_loss = valid_one_epoch(valid_dataloader, model, epoch)\n        \n        train_losses.append(train_loss)\n        valid_losses.append(valid_loss)\n        \n        if valid_loss < best_loss:\n            best_loss = valid_loss\n            save_model(f\"fold{fold}_best_loss.pth\", model, thres)\n            print(\"New Best Loss\")\n        print()\n        \n        print(f\"-------- Epoch {epoch + 1} --------\")\n        print(\"Train Loss: \", train_loss)\n        print(\"Valid Loss: \", valid_loss)\n        print()\n        \n    column_names = ['train_loss','valid_loss']\n    df = pd.DataFrame(np.stack([train_losses, valid_losses], axis=1),columns=column_names)\n    display(df)\n    plot_df(df)","metadata":{"execution":{"iopub.status.busy":"2023-05-24T08:29:17.819170Z","iopub.execute_input":"2023-05-24T08:29:17.819529Z","iopub.status.idle":"2023-05-24T08:29:17.828821Z","shell.execute_reply.started":"2023-05-24T08:29:17.819501Z","shell.execute_reply":"2023-05-24T08:29:17.827767Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for fold in range(CFG.n_folds):\n    if fold not in CFG.train_folds: continue\n    \n    train_df = df[df['fold'] != fold]\n    valid_df = df[df['fold'] == fold]\n    \n    train_dataset = HubmapDataset(train_df, config, transformer(\"train\"))\n    valid_dataset = HubmapDataset(valid_df, config, transformer(\"valid\"))\n    \n    train_dataloader = DataLoader(train_dataset, batch_size=CFG.batch_size, shuffle=True)\n    valid_dataloader = DataLoader(valid_dataset, batch_size=2 * CFG.batch_size, shuffle=False)\n    \n    model = HubmapModel(config).to(device)\n#     model = Mask2FormerForUniversalSegmentation(config).to(device)\n    optimizer, scheduler = get_optimizer_and_scheduler(model, train_dataloader, \"adamw\")\n    \n    train_fnc(train_dataloader, valid_dataloader, model, fold, optimizer, scheduler)","metadata":{"execution":{"iopub.status.busy":"2023-05-24T08:29:18.047370Z","iopub.execute_input":"2023-05-24T08:29:18.047734Z","iopub.status.idle":"2023-05-24T08:29:20.204899Z","shell.execute_reply.started":"2023-05-24T08:29:18.047705Z","shell.execute_reply":"2023-05-24T08:29:20.203566Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}