{"cells":[{"metadata":{"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"!pip install segmentation_models_pytorch","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport os\nimport os\nimport cv2\nimport numpy as np\nimport pandas as pd\nfrom torch.utils.data import TensorDataset, DataLoader, Dataset\nimport albumentations as albu\nimport segmentation_models_pytorch as smp\nimport torch\nfrom tqdm.auto import tqdm","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# data preprocessing - https://www.kaggle.com/iafoss/256x256-images"},{"metadata":{"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"!mkdir data\n!mkdir data/images\n!unzip ../input/512x512-images/train.zip -d data/images","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"!mkdir data/masks\n!unzip ../input/512x512-images/masks.zip -d data/masks","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class config:\n    images_path = './data/images'\n    masks_path = './data/masks'\n    backbone = 'resnet34'\n    lr=1e-3\n    epochs = 10\n    batch_size=8\n    T_max=500\n    im_size=512\n    num_workers=4\n    DEBUG = True #change to False to run all","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def to_tensor(x, **kwargs):\n    return x.transpose(2, 0, 1).astype('float32')\n\ndef get_train_augmentation(size=1024):\n    return albu.Compose([\n        albu.HorizontalFlip(),\n        albu.OneOf([\n            albu.RandomContrast(),\n            albu.RandomGamma(),\n            albu.RandomBrightness(),\n            ], p=0.3),\n        albu.OneOf([\n            albu.ElasticTransform(alpha=120, sigma=120 * 0.05, alpha_affine=120 * 0.03),\n            albu.GridDistortion(),\n            albu.OpticalDistortion(distort_limit=2, shift_limit=0.5),\n            ], p=0.3),\n        albu.ShiftScaleRotate(),\n        albu.Resize(size,size,always_apply=True),\n    ])\n\ndef get_valid_augmentation(size=1024):\n    return albu.Compose([\n        albu.Resize(size,size,always_apply=True),\n    ])\n\n\ndef get_preprocessing(preprocessing_fn):\n    _transform = [\n        albu.Lambda(image=preprocessing_fn),\n        albu.Lambda(image=to_tensor, mask=to_tensor),\n    ]\n    return albu.Compose(_transform)\n\n\n\nclass HuBMAPDataset(Dataset):\n    def __init__(self, ids, transforms=None, preprocessing=None):\n        self.ids = ids\n        self.transforms = transforms\n        self.preprocessing = preprocessing\n    def __getitem__(self, idx):\n        name = self.ids[idx]\n        img = cv2.imread(f\"{config.images_path}/{name}\")\n        mask = cv2.imread(f\"{config.masks_path}/{name}\")[:,:,0:1]\n        if self.transforms:\n            augmented = self.transforms(image=img, mask=mask)\n            img = augmented['image']\n            mask = augmented['mask']\n        if self.preprocessing:\n            preprocessed = self.preprocessing(image=img, mask=mask)\n            img = preprocessed['image']\n            mask = preprocessed['mask']\n        return img, mask\n\n    def __len__(self):\n        return len(self.ids)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"data = os.listdir(config.images_path)\ntrain_lsit = list(set([row.split(\"_\")[0] for row in data]))\ntrain_idx = [row for row in data if row.split(\"_\")[0] in train_lsit[:-2]]\nvalid_idx = [row for row in data if row.split(\"_\")[0] not in train_lsit[:-2]]\nlen(train_idx),len(valid_idx)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"preprocessing = get_preprocessing(smp.encoders.get_preprocessing_fn(\"resnet34\",\"imagenet\"))\ntrain_aug = get_train_augmentation(config.im_size)\nval_aug = get_valid_augmentation(config.im_size)\ntrain_datasets = HuBMAPDataset(train_idx,transforms=train_aug, preprocessing=preprocessing)\nvalid_datasets = HuBMAPDataset(valid_idx,transforms=val_aug, preprocessing=preprocessing)\ntrain_loader = DataLoader(train_datasets, batch_size=config.batch_size, shuffle=True, num_workers=config.num_workers,pin_memory=True)\nvalid_loader = DataLoader(valid_datasets, batch_size=config.batch_size, shuffle=False, num_workers=config.num_workers, pin_memory=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"x,y = train_datasets[1]\nx.shape,y.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = smp.Unet(config.backbone,in_channels = 3,classes = 1,decoder_use_batchnorm = False)\noptim = torch.optim.AdamW(model.parameters(),lr=config.lr)\nscheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optim,T_max=config.T_max)\nloss_fn = torch.nn.BCEWithLogitsLoss()\nmetric = smp.utils.losses.DiceLoss()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class AverageMeter(object):\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\n\nclass trainer:\n    def __init__(self, model,optim,scheduler,loss_fn,metric):\n        self.model = model.cuda()\n        self.opt = optim\n        self.scheduler = scheduler\n        self.loss_fn = loss_fn\n        self.metric = metric\n\n    def train(self, train_loader,e,epochs):\n        self.model.train()\n        tqdm_loder = tqdm(train_loader)\n        current_loss_mean = 0\n        current_dice_mean = AverageMeter()\n        self.opt.zero_grad()\n        for batch_idx, (x, y) in enumerate(tqdm_loder):\n            x = x.cuda().float()\n            y = y.cuda().float()\n            predicted = self.model(x.cuda().float())\n            loss = self.loss_fn(predicted.cuda().float(), y.cuda().float())\n            dice = self.metric(predicted.cuda().float(), y.cuda().float())\n            current_dice_mean.update(dice)\n            predicted_sigmoid = torch.sigmoid(predicted)\n            loss.backward()\n            self.opt.step()\n            self.opt.zero_grad()\n            self.scheduler.step()\n            current_loss_mean = (current_loss_mean * batch_idx + loss) / (batch_idx + 1)\n            lr = self.opt.param_groups[0]['lr']\n            tqdm_loder.set_description(f\"Epoch {e}/{epochs}, train loss: {current_loss_mean:.4} dice {current_dice_mean.avg:.4},lr: {lr:.4}\")\n            if config.DEBUG and batch_idx>10:\n                break\n\n    def valid(self, val_loader,epoch):\n        self.model.eval()\n        tqdm_loder = tqdm(val_loader)\n        current_loss_mean = 0\n        current_dice_mean = AverageMeter()\n        for batch_idx, (x, y) in enumerate(tqdm_loder):\n            with torch.no_grad():\n                x = x.cuda().float()\n                y = y.cuda().float()\n                predicted = self.model(x)\n                predicted_sigmoid = torch.sigmoid(predicted)\n                loss = self.loss_fn(predicted.float(), y.float())\n                dice = self.metric(predicted.cuda().float(), y.cuda().float())\n                current_dice_mean.update(dice)\n            current_loss_mean = (current_loss_mean * batch_idx + loss) / (batch_idx + 1)\n            tqdm_loder.set_description(f\"val loss: {current_loss_mean:.4}, dice {current_dice_mean.avg:.4}\")\n            if config.DEBUG and batch_idx>10:\n                break\n        return current_loss_mean\n    def run(self, train_loader, val_loader,epochs):\n        best = 10000\n        for e in range(epochs):\n            self.train(train_loader,e,epochs)\n            score = self.valid(val_loader,e)\n            if score < best:\n                best=score\n                torch.save(model.state_dict(),\"best.pth\")\n                print(\"save best model\")\n            if config.DEBUG:\n                break\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"epochs = config.epochs\nif config.DEBUG: print(\"DEBUG mode\")\nT = trainer(model,optim,scheduler,loss_fn,metric)\nT.run(train_loader, valid_loader,epochs)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!rm -rf *","execution_count":null,"outputs":[]}],"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}