{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Features\n+ Training with pytorch lightning and a 8 core tpu\n+ Multi-label stratification split\n+ Resnet50 with focal loss and rgby channels\n+ Use jpg images to speed up training speed "},{"metadata":{},"cell_type":"markdown","source":"## Install Deps"},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"!pip install iterative-stratification\n!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n!python pytorch-xla-env-setup.py --version 1.7 --apt-packages libomp5 libopenblas-dev\n!pip install -U pytorch-lightning==1.1.4 albumentations","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Import packages"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import pandas as pd\nfrom PIL import Image\nimport os\nimport numpy as np\nfrom io import BytesIO\n\nfrom matplotlib import pyplot as plt\n\nfrom tqdm.notebook import tqdm\nimport traceback\n\nimport pandas as pd\nimport numpy as np\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import DataLoader, Dataset\nfrom torchvision import models\n\nimport cv2\n\nfrom pytorch_lightning import LightningModule, Trainer, seed_everything\nfrom pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor\nfrom pytorch_lightning.profiler import AdvancedProfiler\nfrom pytorch_lightning.loggers import TensorBoardLogger\nimport time\nimport sys","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Hyper params"},{"metadata":{"trusted":true},"cell_type":"code","source":"img_w, img_h = 512, 512\nn_classes = 19\n\n\nbatch_size = 8\nepoch = 1\nlr = 3e-4","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Image Loading Pipeline"},{"metadata":{"trusted":true},"cell_type":"code","source":"df = pd.read_csv('../input/hpa-single-cell-image-classification/train.csv')\ndf","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\nfrom torchvision import transforms\n\naug_train = A.Compose([\n    A.Resize(img_h, img_w, interpolation=cv2.INTER_LINEAR),\n    A.HorizontalFlip(p=0.5),\n    A.RandomRotate90(),\n    A.Normalize(mean=(0.5, 0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5, 0.5)),\n    ToTensorV2(transpose_mask=True)\n])\n\n\naug_val = A.Compose([\n    A.Resize(img_h, img_w, interpolation=cv2.INTER_LINEAR),\n    A.Normalize(mean=(0.5, 0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5, 0.5)),\n    ToTensorV2(transpose_mask=True)\n])\n\nclass MyDataset(Dataset):\n    def __init__(self, df, aug):\n        self.imgs = df['ID'].values\n        self.labels = df['Label'].values\n        self.aug = aug\n        \n    def __len__(self):\n        return len(self.imgs)\n    \n    def __getitem__(self, idx):\n        img = cv2.imread(f'../input/hpa-image-loading-speed-up-for-training/train/{self.imgs[idx]}_rgb.jpg', cv2.IMREAD_UNCHANGED)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        imgy = cv2.imread(f'../input/hpa-image-loading-speed-up-for-training/train/{self.imgs[idx]}_yellow.jpg', cv2.IMREAD_UNCHANGED)\n        \n        img = np.concatenate((img, imgy.reshape((imgy.shape[0], imgy.shape[1], 1))), axis=2)\n        \n        aug_rst = self.aug(image=img)\n        img = aug_rst['image']\n        \n        label = torch.zeros(n_classes, dtype=torch.long)\n        \n        for x in self.labels[idx].split('|'):\n            x = int(x)\n            \n            label[x] = 1\n        \n        return img, label\n\n\ndef reverse_transform(a):\n    r = a[0].numpy()\n    g = a[1].numpy()\n    b = a[2].numpy()\n    y = a[3].numpy()\n    \n    return (np.dstack([r * 0.5 + 0.5, g * 0.5 + 0.5, b * 0.5 + 0.5]) * 255).astype(np.uint8)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import random\n\nds = MyDataset(df, aug_train)\n\nplt.figure(figsize=(20, 8))\nfor i in range(10):\n    idx = random.randrange(0, len(ds))\n    d = ds[idx]\n    \n    plt.subplot(2, 5, i + 1)\n    plt.imshow(reverse_transform(d[0]))\n    plt.title(df.iloc[idx].Label)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Define Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"sys.path.append('../input/pretrainedmodels')\nimport pretrainedmodels\n\n!mkdir -p /root/.cache/torch/hub/checkpoints/\n!ln -s /kaggle/input/pretrained-model-weights-pytorch/* /root/.cache/torch/hub/checkpoints/","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class FocalLoss(nn.Module):\n    def __init__(self, gamma=2):\n        super().__init__()\n        \n        self.gamma = gamma\n        \n        self.bce = nn.BCEWithLogitsLoss(reduction='none')\n        \n    def forward(self, logits, targets):\n        logpt = self.bce(logits, targets)\n        \n        pt = torch.exp(-logpt)\n        \n        loss = ((1 - pt) ** self.gamma) * logpt\n        \n        return loss.mean()\n\n\nclass Model(LightningModule):\n    def __init__(self, lr, batch_size, steps_per_epoch, n_epoch):\n        super().__init__()\n        \n        self.save_hyperparameters()\n        self.lr = lr\n        \n        self.model = pretrainedmodels.resnet50(pretrained='imagenet')\n        \n        stem = nn.Conv2d(4, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n        w = self.model.conv1.weight\n        stem.weight = nn.Parameter(torch.cat((w, 0.5*(w[:,:1,:,:]+w[:,2:,:,:])),dim=1))\n        self.model.conv1 = stem\n\n        self.avg_pool = nn.AdaptiveAvgPool2d((1, 1))\n        self.fc = nn.Linear(2048, n_classes)\n        \n        self.crit = FocalLoss()\n\n        \n    def forward(self, ipt):\n        x = self.model.features(ipt)\n        x = self.avg_pool(x)\n        x = self.fc(x.view(-1, 2048))\n        return x\n    \n    def training_step(self, batch, batch_idx):\n        self.train()\n        x, y = batch\n        \n        pred = self(x)\n        \n        loss = self.crit(pred, y.float())\n\n        rst = (pred > 0.5).long()\n        acc = (rst == y).float().mean()\n        \n        self.log('train_loss', loss)\n        self.log('train_acc', acc)\n        \n        return loss\n    \n    def configure_optimizers(self):\n        optimizer = torch.optim.AdamW(self.parameters(), lr=self.lr, weight_decay=1e-4)\n        scheduler = {\n            'scheduler': torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=self.lr, \n                                                             total_steps=self.hparams.n_epoch * self.hparams.steps_per_epoch,\n                                                             anneal_strategy='cos',\n                                                             cycle_momentum=False,\n                                                             pct_start=0.1,\n                                                            ),\n            'interval': 'step',\n            'frequency': 1\n        }\n        \n        return [optimizer], [scheduler]\n    \n    \n    def validation_step(self, batch, batch_idx):\n        with torch.no_grad():\n            self.eval()\n            x, y = batch\n            \n            pred = self(x)\n        \n            loss = self.crit(pred, y.float())\n\n            rst = (pred > 0.5).long()\n            acc = (rst == y).float().mean()\n            \n            self.log('val_loss', loss)\n            self.log('val_acc', acc)\n\ndef init_seed(worker_id):\n    random.seed(torch.torch.initial_seed())","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Calc Label for MultilabelStratifiedKFold"},{"metadata":{"trusted":true},"cell_type":"code","source":"Y = []\n\nfor row in tqdm(df.itertuples()):\n    y = np.zeros(n_classes, dtype=np.int64)\n    \n    for _id in row.Label.split('|'):\n        y[int(_id)] = 1\n        \n    Y.append(y)\n    \nY = np.stack(Y)\nY","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Start Training"},{"metadata":{"trusted":true},"cell_type":"code","source":"from sklearn.model_selection import train_test_split, StratifiedKFold, KFold\nfrom iterstrat.ml_stratifiers import MultilabelStratifiedKFold\n\n\nskf = MultilabelStratifiedKFold(n_splits=5, random_state=42, shuffle=True)\n\nfor fold_idx, (train, val) in enumerate(skf.split(df, Y)):\n    print('training fold', fold_idx)\n    df_train = df.loc[train].sort_index()\n    df_val = df.loc[val].sort_index()\n\n    ds_train = MyDataset(df_train, aug_train)\n\n    print('train samples:', len(df_train))\n    train_loader = DataLoader(ds_train, batch_size=batch_size, num_workers=0, shuffle=True,\n                              worker_init_fn=init_seed, drop_last=True)\n\n    ds_val = MyDataset(df_val, aug_val)\n    val_loader = DataLoader(ds_val, batch_size=batch_size, num_workers=0, drop_last=True)\n\n    checkpoint_callback = ModelCheckpoint(\n        save_top_k=1,\n        save_last=True,\n        verbose=True,\n        monitor='val_acc',\n        mode='max',\n        prefix=''\n    )\n\n    logger = TensorBoardLogger(save_dir='.', version=f'fold_{fold_idx}', name='lightning_logs')\n\n    trainer = Trainer(logger=logger,\n                      tpu_cores=8,\n                      callbacks=[LearningRateMonitor(), checkpoint_callback],\n                      flush_logs_every_n_steps=100,\n                      log_every_n_steps=100,\n                      max_epochs=epoch,\n                      benchmark=True,\n                      precision=16)\n\n    model = Model(lr, batch_size, len(train_loader), epoch)\n    trainer.fit(model, train_dataloader=train_loader, val_dataloaders=val_loader)\n\n    break","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Inference\n1. Obtain cell instance masks using hpaCellSegmentator\n2. Mask the input image to get single cell images\n3. Run classification model with single cell images to get the final result"}],"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":4,"nbformat_minor":4}