{"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":"code","source":"from sklearn.metrics import roc_auc_score, accuracy_score, f1_score, log_loss\nimport pickle\nfrom torch.utils.data import DataLoader\nfrom torch.cuda.amp import autocast, GradScaler\nimport warnings\nimport sys\nimport pandas as pd\nimport os\nimport gc\nimport sys\nimport math\nimport time\nimport random\nimport shutil\nfrom pathlib import Path\nfrom contextlib import contextmanager\nfrom collections import defaultdict, Counter\nimport cv2\n\nimport scipy as sp\nimport numpy as np\nimport pandas as pd\n\nimport matplotlib.pyplot as plt\nfrom tqdm.auto import tqdm\nfrom functools import partial\n\nimport argparse\nimport importlib\nimport torch\nimport torch.nn as nn\nfrom torch.optim import Adam, SGD, AdamW\n\nimport datetime\nimport wandb","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install segmentation_models_pytorch\n!pip install warmup_scheduler","metadata":{"collapsed":true,"jupyter":{"outputs_hidden":true},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import segmentation_models_pytorch as smp","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nfrom torch.utils.data import DataLoader, Dataset\nimport cv2\nimport torch\nimport os\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nfrom albumentations import ImageOnlyTransform","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## config","metadata":{}},{"cell_type":"code","source":"import os\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\nclass CFG:\n    # ============== comp exp name =============\n    comp_name = 'vesuvius'\n    exp_name = 'mit_b2'\n    comp_dir_path = '/kaggle/input/'\n    comp_folder_name = 'vesuvius-png'\n    comp_dataset_path = f'{comp_dir_path}{comp_folder_name}/train/'\n\n#     comp_dir_path = '/kaggle/input/'\n#     comp_folder_name = 'vesuvius-challenge-ink-detection'\n#     comp_dataset_path = f'{comp_dir_path}{comp_folder_name}/'\n\n    \n    exp_name = 'vesuvius_2d_slide_exp002'\n\n    # ============== pred target =============\n    target_size = 1\n\n    # ============== model cfg =============\n    model_name = 'Unet'\n    backbone = 'mit_b2'\n\n    in_chans = 3 # 65\n    # ============== training cfg =============\n    size = 224\n    tile_size = 224\n    stride = tile_size // 4\n\n    train_batch_size = 128 # 32\n    valid_batch_size = train_batch_size * 1\n    use_amp = True\n\n    scheduler = 'GradualWarmupSchedulerV2'\n    # scheduler = 'CosineAnnealingLR'\n    epochs = 5\n    # 30\n\n    warmup_factor = 10\n    lr = 1e-4 / warmup_factor\n    # adamW warmupあり\n\n    # ============== fold =============\n    valid_id = [3]\n\n    # objective_cv = 'binary'  # 'binary', 'multiclass', 'regression'\n    metric_direction = 'maximize'  # maximize, 'minimize'\n    # metrics = 'dice_coef'\n\n    # ============== fixed =============\n    pretrained = True\n    inf_weight = 'best'  # 'best'\n\n    min_lr = 1e-6\n    weight_decay = 1e-6\n    max_grad_norm = 1000\n\n    print_freq = 50\n    num_workers = 2\n\n    seed = 42\n\n    # ============== set dataset path =============\n    print('set dataset path')\n\n    outputs_path = f'/kaggle/working/outputs/{comp_name}/{exp_name}/'\n\n    submission_dir = outputs_path + 'submissions/'\n    submission_path = submission_dir + f'submission_{exp_name}.csv'\n\n    model_dir = outputs_path + \\\n        f'{comp_name}-models/'\n\n    figures_dir = outputs_path + 'figures/'\n\n    log_dir = outputs_path + 'logs/'\n    log_path = log_dir + f'{exp_name}.txt'\n\n    # ============== augmentation =============\n    train_aug_list = [\n        # A.RandomResizedCrop(\n        #     size, size, scale=(0.85, 1.0)),\n        A.Resize(size, size),\n        A.HorizontalFlip(p=0.5),\n        A.VerticalFlip(p=0.5),\n        A.ShiftScaleRotate(p=0.5),\n        A.OneOf([\n            A.ElasticTransform(p=.3),\n            A.GaussianBlur(p=.3),\n            A.GaussNoise(p=.3),\n            A.OpticalDistortion(p=0.3),\n            A.GridDistortion(p=.1),\n        ], p=0.3),\n        A.Normalize(\n            mean= [0] * in_chans,\n            std= [1] * in_chans\n        ),\n        ToTensorV2(transpose_mask=True),\n    ]\n\n    valid_aug_list = [\n        A.Resize(size, size),\n        A.Normalize(\n            mean= [0] * in_chans,\n            std= [1] * in_chans\n        ),\n        ToTensorV2(transpose_mask=True),\n    ]\n\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## helper","metadata":{}},{"cell_type":"code","source":"class AverageMeter(object):\n    \"\"\"Computes and stores the average and current value\"\"\"\n\n    def __init__(self):\n        self.reset()\n\n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n\n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def init_logger(log_file):\n    from logging import getLogger, INFO, FileHandler, Formatter, StreamHandler\n    logger = getLogger(__name__)\n    logger.setLevel(INFO)\n    handler1 = StreamHandler()\n    handler1.setFormatter(Formatter(\"%(message)s\"))\n    handler2 = FileHandler(filename=log_file)\n    handler2.setFormatter(Formatter(\"%(message)s\"))\n    logger.addHandler(handler1)\n    logger.addHandler(handler2)\n    return logger\n\ndef set_seed(seed=None, cudnn_deterministic=True):\n    if seed is None:\n        seed = 42\n\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = cudnn_deterministic\n    torch.backends.cudnn.benchmark = False","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def make_dirs(cfg):\n    for dir in [cfg.model_dir, cfg.figures_dir, cfg.submission_dir, cfg.log_dir]:\n        os.makedirs(dir, exist_ok=True)\ndef cfg_init(cfg, mode='train'):\n    set_seed(cfg.seed)\n    # set_env_name()\n    # set_dataset_path(cfg)\n\n    if mode == 'train':\n        make_dirs(cfg)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cfg_init(CFG)\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\nLogger = init_logger(log_file=CFG.log_path)\n\nLogger.info('\\n\\n-------- exp_info -----------------')\n# Logger.info(datetime.datetime.now().strftime('%Y年%m月%d日 %H:%M:%S'))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## image, mask","metadata":{}},{"cell_type":"code","source":"def read_image_mask(fragment_id):\n\n    images = []\n    \n    mid = 65 // 2\n    start = 29\n    end = 32\n    idxs = range(start, end)\n\n    for i in tqdm(idxs):\n        \n        image = cv2.imread(CFG.comp_dataset_path + f\"train/{fragment_id}/surface_volume/{i:02}.png\", 0)\n\n        pad0 = (CFG.tile_size - image.shape[0] % CFG.tile_size)\n        pad1 = (CFG.tile_size - image.shape[1] % CFG.tile_size)\n\n        image = np.pad(image, [(0, pad0), (0, pad1)], constant_values=0)\n\n        images.append(image)\n    images = np.stack(images, axis=2)\n\n    mask = cv2.imread(CFG.comp_dataset_path + f\"train/{fragment_id}/inklabels.png\", 0)\n    mask = np.pad(mask, [(0, pad0), (0, pad1)], constant_values=0)\n\n    mask = mask.astype('float32')\n    mask /= 255.0\n    \n    return images, mask","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_train_valid_dataset(valid_id):\n    train_images = []\n    train_masks = []\n\n    valid_images = []\n    valid_masks = []\n    valid_xyxys = []\n\n    for fragment_id in range(1, 4):\n\n        image, mask = read_image_mask(fragment_id)\n\n        x1_list = list(range(0, image.shape[1]-CFG.tile_size+1, CFG.stride))\n        y1_list = list(range(0, image.shape[0]-CFG.tile_size+1, CFG.stride))\n\n        for y1 in y1_list:\n            for x1 in x1_list:\n                y2 = y1 + CFG.tile_size\n                x2 = x1 + CFG.tile_size\n                # xyxys.append((x1, y1, x2, y2))\n        \n                if fragment_id == valid_id:\n                    valid_images.append(image[y1:y2, x1:x2])\n                    valid_masks.append(mask[y1:y2, x1:x2, None])\n\n                    valid_xyxys.append([x1, y1, x2, y2])\n                else:\n                    train_images.append(image[y1:y2, x1:x2])\n                    train_masks.append(mask[y1:y2, x1:x2, None])\n\n    return train_images, train_masks, valid_images, valid_masks, valid_xyxys","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## dataset","metadata":{}},{"cell_type":"code","source":"import numpy as np\nfrom torch.utils.data import DataLoader, Dataset\nimport cv2\nimport torch\nimport os\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nfrom albumentations import ImageOnlyTransform","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_transforms(data, cfg):\n    if data == 'train':\n        aug = A.Compose(cfg.train_aug_list)\n    elif data == 'valid':\n        aug = A.Compose(cfg.valid_aug_list)\n\n    # print(aug)\n    return aug\n\nclass CustomDataset(Dataset):\n    def __init__(self, images, cfg, labels=None, transform=None):\n        self.images = images\n        self.cfg = cfg\n        self.labels = labels\n        self.transform = transform\n\n    def __len__(self):\n        # return len(self.df)\n        return len(self.images)\n\n    def __getitem__(self, idx):\n        image = self.images[idx]\n        label = self.labels[idx]\n\n        if self.transform:\n            data = self.transform(image=image, mask=label)\n            image = data['image']\n            label = data['mask']\n\n        return image, label","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CustomModel(nn.Module):\n    def __init__(self, cfg, weight=None):\n        super().__init__()\n        self.cfg = cfg\n\n        self.encoder = smp.Unet(\n            encoder_name=cfg.backbone, \n            encoder_weights=weight,\n            in_channels=cfg.in_chans,\n            classes=cfg.target_size,\n            activation=None,\n            decoder_attention_type = 'scse'\n            \n        )\n\n    def forward(self, image):\n        output = self.encoder(image)\n        # output = output.squeeze(-1)\n        return output\n    \ndef build_model(cfg, weight=\"imagenet\"):\n    print('model_name', cfg.model_name)\n    print('backbone', cfg.backbone)\n\n    model = CustomModel(cfg, weight)\n\n    return model","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## model","metadata":{}},{"cell_type":"markdown","source":"## scheduler","metadata":{}},{"cell_type":"code","source":"\nimport torch.nn as nn\nimport torch\nimport math\nimport time\nimport numpy as np\nimport torch\n\nfrom torch.optim.lr_scheduler import CosineAnnealingWarmRestarts, CosineAnnealingLR, ReduceLROnPlateau\nfrom warmup_scheduler import GradualWarmupScheduler\n\n\nclass GradualWarmupSchedulerV2(GradualWarmupScheduler):\n    \"\"\"\n    https://www.kaggle.com/code/underwearfitting/single-fold-training-of-resnet200d-lb0-965\n    \"\"\"\n    def __init__(self, optimizer, multiplier, total_epoch, after_scheduler=None):\n        super(GradualWarmupSchedulerV2, self).__init__(\n            optimizer, multiplier, total_epoch, after_scheduler)\n\n    def get_lr(self):\n        if self.last_epoch > self.total_epoch:\n            if self.after_scheduler:\n                if not self.finished:\n                    self.after_scheduler.base_lrs = [\n                        base_lr * self.multiplier for base_lr in self.base_lrs]\n                    self.finished = True\n                return self.after_scheduler.get_lr()\n            return [base_lr * self.multiplier for base_lr in self.base_lrs]\n        if self.multiplier == 1.0:\n            return [base_lr * (float(self.last_epoch) / self.total_epoch) for base_lr in self.base_lrs]\n        else:\n            return [base_lr * ((self.multiplier - 1.) * self.last_epoch / self.total_epoch + 1.) for base_lr in self.base_lrs]\n\ndef get_scheduler(cfg, optimizer):\n    scheduler_cosine = torch.optim.lr_scheduler.CosineAnnealingLR(\n        optimizer, cfg.epochs, eta_min=1e-7)\n    scheduler = GradualWarmupSchedulerV2(\n        optimizer, multiplier=10, total_epoch=1, after_scheduler=scheduler_cosine)\n\n    return scheduler\n\ndef scheduler_step(scheduler, avg_val_loss, epoch):\n    scheduler.step(epoch)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## loss","metadata":{}},{"cell_type":"code","source":"DiceLoss = smp.losses.DiceLoss(mode='binary')\nBCELoss = smp.losses.SoftBCEWithLogitsLoss()\nLovaszLoss = smp.losses.LovaszLoss(mode='binary')\nJaccardLoss = smp.losses.JaccardLoss(mode='binary')\ndef criterion(y_pred, y_true):\n     return 0.5*BCELoss(y_pred, y_true) + 0.5 * DiceLoss(y_pred, y_true)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## train, val","metadata":{}},{"cell_type":"code","source":"def train_fn(train_loader, model, criterion, optimizer, device):\n    model.train()\n\n    scaler = GradScaler(enabled=CFG.use_amp)\n    losses = AverageMeter()\n\n    for step, (images, labels) in tqdm(enumerate(train_loader), total=len(train_loader)):\n        images = images.to(device)\n        labels = labels.to(device)\n        batch_size = labels.size(0)\n\n        with autocast(CFG.use_amp):\n            y_preds = model(images)\n            loss = criterion(y_preds, labels)\n\n        losses.update(loss.item(), batch_size)\n        scaler.scale(loss).backward()\n\n        grad_norm = torch.nn.utils.clip_grad_norm_(\n            model.parameters(), CFG.max_grad_norm)\n\n        scaler.step(optimizer)\n        scaler.update()\n        optimizer.zero_grad()\n\n    return losses.avg\n\ndef valid_fn(valid_loader, model, criterion, device, valid_xyxys, valid_mask_gt):\n    mask_pred = np.zeros(valid_mask_gt.shape)\n    mask_count = np.zeros(valid_mask_gt.shape)\n\n    model.eval()\n    losses = AverageMeter()\n\n    for step, (images, labels) in tqdm(enumerate(valid_loader), total=len(valid_loader)):\n        images = images.to(device)\n        labels = labels.to(device)\n        batch_size = labels.size(0)\n\n        with torch.no_grad():\n            y_preds = model(images)\n            loss = criterion(y_preds, labels)\n        losses.update(loss.item(), batch_size)\n\n        # make whole mask\n        y_preds = torch.sigmoid(y_preds).to('cpu').numpy()\n        start_idx = step*CFG.valid_batch_size\n        end_idx = start_idx + batch_size\n        for i, (x1, y1, x2, y2) in enumerate(valid_xyxys[start_idx:end_idx]):\n            mask_pred[y1:y2, x1:x2] += y_preds[i].squeeze(0)\n            mask_count[y1:y2, x1:x2] += np.ones((CFG.tile_size, CFG.tile_size))\n\n    print(f'mask_count_min: {mask_count.min()}')\n    mask_pred /= mask_count\n    return losses.avg, mask_pred","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## metrics","metadata":{}},{"cell_type":"code","source":"from sklearn.metrics import fbeta_score\n\n\n\ndef fbeta_numpy(targets, preds, beta=0.5, smooth=1e-5):\n    \"\"\"\n    https://www.kaggle.com/competitions/vesuvius-challenge-ink-detection/discussion/397288\n    \"\"\"\n    y_true_count = targets.sum()\n    ctp = preds[targets==1].sum()\n    cfp = preds[targets==0].sum()\n    beta_squared = beta * beta\n\n    c_precision = ctp / (ctp + cfp + smooth)\n    c_recall = ctp / (y_true_count + smooth)\n    dice = (1 + beta_squared) * (c_precision * c_recall) / (beta_squared * c_precision + c_recall + smooth)\n\n    return dice\n\ndef calc_fbeta(mask, mask_pred, _range=(20,80+1), _step=5):\n    mask = mask.astype(int).flatten()\n    mask_pred = mask_pred.flatten()\n\n    best_th = 0\n    best_dice = 0\n    for th in np.array(range(*_range, _step)) / 100:\n        \n        # dice = fbeta_score(mask, (mask_pred >= th).astype(int), beta=0.5)\n        dice = fbeta_numpy(mask, (mask_pred >= th).astype(int), beta=0.5)\n        print(f'th: {th}, fbeta: {dice}')\n\n        if dice > best_dice:\n            best_dice = dice\n            best_th = th\n        else:\n            break\n    \n    Logger.info(f'best_th: {best_th}, fbeta: {best_dice}')\n    torch.cuda.empty_cache()\n    gc.collect()\n    return best_dice, best_th\n\n\ndef calc_cv(mask_gt, mask_pred, _range=(20,80+1), prec=5):\n    best_dice, best_th = calc_fbeta(mask_gt, mask_pred, _range=_range,_step=prec)\n\n    return best_dice, best_th","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## main","metadata":{}},{"cell_type":"code","source":"for fold in CFG.valid_id:\n    \n    train_images, train_masks, valid_images, valid_masks, valid_xyxys = get_train_valid_dataset(fold)\n    valid_xyxys = np.stack(valid_xyxys)\n    \n    \n    train_dataset = CustomDataset(train_images, CFG, labels=train_masks,\n                                  transform=get_transforms(data='train', cfg=CFG))\n\n    valid_dataset = CustomDataset(valid_images, CFG, labels=valid_masks,\n                                  transform=get_transforms(data='valid', cfg=CFG))\n\n    train_loader = DataLoader(train_dataset,\n                              batch_size=CFG.train_batch_size,\n                              shuffle=True,\n                              num_workers=CFG.num_workers, pin_memory=True, drop_last=True,\n                              )\n\n    valid_loader = DataLoader(valid_dataset,\n                              batch_size=CFG.valid_batch_size,\n                              shuffle=False,\n                              num_workers=CFG.num_workers, pin_memory=True, drop_last=False)\n    \n    \n    valid_mask_gt = cv2.imread(CFG.comp_dataset_path + f\"train/{fold}/inklabels.png\", 0)\n    valid_mask_gt = valid_mask_gt / 255\n    pad0 = (CFG.tile_size - valid_mask_gt.shape[0] % CFG.tile_size)\n    pad1 = (CFG.tile_size - valid_mask_gt.shape[1] % CFG.tile_size)\n    valid_mask_gt = np.pad(valid_mask_gt, [(0, pad0), (0, pad1)], constant_values=0)\n    \n    \n    if CFG.metric_direction == 'minimize':\n        best_score = np.inf\n    elif CFG.metric_direction == 'maximize':\n        best_score = -1\n\n    best_loss = np.inf\n    \n    model = build_model(CFG)\n    model.to(device)\n    model = nn.DataParallel(model, device_ids=[0, 1])\n    optimizer = AdamW(model.parameters(), lr=CFG.lr)\n    scheduler = get_scheduler(CFG, optimizer)\n\n    for epoch in range(CFG.epochs):\n\n        start_time = time.time()\n\n        # train\n        avg_loss = train_fn(train_loader, model, criterion, optimizer, device)\n\n        # eval\n        avg_val_loss, mask_pred = valid_fn(\n            valid_loader, model, criterion, device, valid_xyxys, valid_mask_gt)\n\n        scheduler_step(scheduler, avg_val_loss, epoch)\n#         scheduler.step()\n\n        best_dice, best_th = calc_cv(valid_mask_gt, mask_pred)\n\n        # score = avg_val_loss\n        score = best_dice\n\n        elapsed = time.time() - start_time\n\n        Logger.info(\n            f'Epoch {epoch+1} - avg_train_loss: {avg_loss:.4f}  avg_val_loss: {avg_val_loss:.4f}  time: {elapsed:.0f}s')\n        # Logger.info(f'Epoch {epoch+1} - avgScore: {avg_score:.4f}')\n        Logger.info(\n            f'Epoch {epoch+1} - avgScore: {score:.4f}')\n\n        if CFG.metric_direction == 'minimize':\n            update_best = score < best_score\n        elif CFG.metric_direction == 'maximize':\n            update_best = score > best_score\n\n        if update_best:\n            best_loss = avg_val_loss\n            best_score = score\n\n            Logger.info(\n                f'Epoch {epoch+1} - Save Best Score: {best_score:.4f} Model')\n            Logger.info(\n                f'Epoch {epoch+1} - Save Best Loss: {best_loss:.4f} Model')\n\n            torch.save({'model': model.state_dict()},\n                        CFG.model_dir + f'{CFG.model_name}_fold{fold}_best.pth')\n        torch.cuda.empty_cache()\n        gc.collect()\n    del model\n    gc.collect()\n    del optimizer\n    del scheduler\n    gc.collect()\n    torch.cuda.empty_cache()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"best_dice, best_th  = calc_fbeta(valid_mask_gt, mask_pred)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fig, axes = plt.subplots(1, 3, figsize=(15, 8))\naxes[0].imshow(valid_mask_gt)\naxes[1].imshow(mask_pred)\naxes[2].imshow((mask_pred>=best_th).astype(int))","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}