{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Segmentation models\nThis notebook shows how to use the awesome segmentation models package https://github.com/qubvel/segmentation_models.pytorch to create a U-net segmentation model that will try to predict the catheter position given the annotations that accompany this dataset. This can be used as an auxiliary training task or used as downstream input to a different model. This notebook will generate masks for a single validation fold, but can easily be extended to make predictions for full out of fold training set. \n\nThere are several different ways this can be posed, but I have chosen to pose this as a 4 class segmentation so that the model is just trying to identify the 4 different catheter types and not generating masks trying to predict if they are correctly placed or not. It should be trivial to extend this to 11 classes or reduced to 1. \n\nThis model does not train the classification branch but the stub is still there if you would like to change the dataloader and some other parts to train it. "},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install segmentation-models-pytorch","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","execution":{"iopub.execute_input":"2021-01-05T05:34:07.512249Z","iopub.status.busy":"2021-01-05T05:34:07.511443Z","iopub.status.idle":"2021-01-05T05:34:08.145868Z","shell.execute_reply":"2021-01-05T05:34:08.145375Z"},"papermill":{"duration":0.653913,"end_time":"2021-01-05T05:34:08.145979","exception":false,"start_time":"2021-01-05T05:34:07.492066","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"import os\nimage_size = 512\nseed = 42\nuse_amp = True\ndebug = False\nmodel_name = 'resnet18'\nphase = \"only_seg\"\nos.mkdir(\"./models\")\nif os.path.exists(\"./models/\" + model_name):\n    print(\"Warning model directory already exists \\n\"*5)\nelse:\n    os.mkdir(\"./models/\" + model_name)\ndata_dir = '../input/ranzcr-clip-catheter-line-classification/train/'","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","execution":{"iopub.execute_input":"2021-01-05T05:34:20.104275Z","iopub.status.busy":"2021-01-05T05:34:20.103492Z","iopub.status.idle":"2021-01-05T05:34:23.725655Z","shell.execute_reply":"2021-01-05T05:34:23.724300Z"},"papermill":{"duration":3.642867,"end_time":"2021-01-05T05:34:23.725771","exception":false,"start_time":"2021-01-05T05:34:20.082904","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport sys\nsys.path.append('../input/timm-pytorch-image-models/pytorch-image-models-master')\nimport timm\nimport segmentation_models_pytorch as smp\nimport time\nimport cv2\nimport PIL.Image\nimport random\nfrom sklearn.metrics import accuracy_score\nimport torch\nfrom torch.utils.data import DataLoader, Dataset\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision\nimport torchvision.transforms as transforms\nimport torch.optim as optim\nfrom torch.optim.lr_scheduler import CosineAnnealingLR \nimport albumentations as a\nfrom albumentations import *\nfrom albumentations.pytorch import ToTensorV2\nfrom tqdm.notebook import tqdm\nimport matplotlib.pyplot as plt\nimport gc\nfrom sklearn.metrics import roc_auc_score\nimport seaborn as sns\nfrom pylab import rcParams\nimport timm\nimport ast\nfrom warnings import filterwarnings\nfilterwarnings(\"ignore\")\n\ndevice = torch.device('cuda')","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.010976,"end_time":"2021-01-05T05:34:23.748081","exception":false,"start_time":"2021-01-05T05:34:23.737105","status":"completed"},"tags":[]},"cell_type":"markdown","source":"## Model"},{"metadata":{"execution":{"iopub.execute_input":"2021-01-05T05:34:23.780605Z","iopub.status.busy":"2021-01-05T05:34:23.778974Z","iopub.status.idle":"2021-01-05T05:34:23.781320Z","shell.execute_reply":"2021-01-05T05:34:23.781735Z"},"papermill":{"duration":0.022682,"end_time":"2021-01-05T05:34:23.781850","exception":false,"start_time":"2021-01-05T05:34:23.759168","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"class SegNet(nn.Module):\n    def __init__(self, model_name='resnet200d', out_dim=11, pretrained=False):\n        super().__init__()\n        aux_params=dict(\n                        pooling='max',\n                        dropout=0.1,\n                        classes=out_dim)\n        self.model = smp.Unet(model_name, encoder_weights=\"imagenet\", classes = 4, aux_params = aux_params)\n    def forward(self, x):\n        mask_logits, logits = self.model(x)\n        mask_logits = mask_logits.permute(0, 2, 3, 1)\n        return mask_logits, logits\n        \n        ","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.011021,"end_time":"2021-01-05T05:34:23.804023","exception":false,"start_time":"2021-01-05T05:34:23.793002","status":"completed"},"tags":[]},"cell_type":"markdown","source":"## Utils"},{"metadata":{"execution":{"iopub.execute_input":"2021-01-05T05:34:23.837932Z","iopub.status.busy":"2021-01-05T05:34:23.837370Z","iopub.status.idle":"2021-01-05T05:34:23.843545Z","shell.execute_reply":"2021-01-05T05:34:23.843043Z"},"papermill":{"duration":0.028285,"end_time":"2021-01-05T05:34:23.843626","exception":false,"start_time":"2021-01-05T05:34:23.815341","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"def seed_everything(seed):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = True # for faster training, but not deterministic\n    \nseed_everything(seed)","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.011114,"end_time":"2021-01-05T05:34:23.868630","exception":false,"start_time":"2021-01-05T05:34:23.857516","status":"completed"},"tags":[]},"cell_type":"markdown","source":"## Transforms"},{"metadata":{"execution":{"iopub.execute_input":"2021-01-05T05:34:23.903753Z","iopub.status.busy":"2021-01-05T05:34:23.903220Z","iopub.status.idle":"2021-01-05T05:34:23.907369Z","shell.execute_reply":"2021-01-05T05:34:23.906782Z"},"papermill":{"duration":0.027104,"end_time":"2021-01-05T05:34:23.907462","exception":false,"start_time":"2021-01-05T05:34:23.880358","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"transforms_train = a.Compose([\n   a.RandomResizedCrop(image_size, image_size, scale=(0.9, 1), p=1), \n   a.HorizontalFlip(p=0.5),\n   a.ShiftScaleRotate(p=0.5),\n   a.HueSaturationValue(hue_shift_limit=10, sat_shift_limit=10, val_shift_limit=10, p=0.7),\n   a.RandomBrightnessContrast(brightness_limit=(-0.2,0.2), contrast_limit=(-0.2, 0.2), p=0.7),\n   a.CLAHE(clip_limit=(1,4), p=0.5),\n   a.OneOf([\n       a.OpticalDistortion(distort_limit=1.0),\n       a.GridDistortion(num_steps=5, distort_limit=1.),\n       a.ElasticTransform(alpha=3),\n   ], p=0.2),\n   a.OneOf([\n       a.GaussNoise(var_limit=[10, 50]),\n       a.GaussianBlur(),\n       a.MotionBlur(),\n       a.MedianBlur(),\n   ], p=0.2),\n  a.Resize(image_size, image_size),\n  a.OneOf([\n      JpegCompression(),\n      Downscale(scale_min=0.1, scale_max=0.15),\n  ], p=0.2),\n  IAAPiecewiseAffine(p=0.2),\n  IAASharpen(p=0.2),\n  a.Cutout(max_h_size=int(image_size * 0.1), max_w_size=int(image_size * 0.1), num_holes=5, p=0.5),\n  a.Normalize(),\n  ToTensorV2()\n])\n\ntransforms_valid = a.Compose([\n    a.Resize(image_size, image_size),\n    a.Normalize(),\n    ToTensorV2()\n])","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.011525,"end_time":"2021-01-05T05:34:23.930452","exception":false,"start_time":"2021-01-05T05:34:23.918927","status":"completed"},"tags":[]},"cell_type":"markdown","source":"## Dataset"},{"metadata":{},"cell_type":"markdown","source":"This interpolation function is used to extend the annotations so we have the smooth curve connecting the points given in the annotation. In the annotation labels we are given a series of points that trace the path of catheter/line, but it would be hard for a model to predict these points because the spacing in between them is a bit arbitrary. To remedy this a straight line is drawn between each consecutive point. It is not perfect, but yields reasonably good looking outputs. This is what the model will be trying to predict. "},{"metadata":{"trusted":true},"cell_type":"code","source":"from scipy import interpolate\ndef interpolate_mask(data):\n    f = interpolate.interp1d(data[:, 0], data[:, 1])\n    xnew = np.arange(data[:, 0].min(), data[:, 0].max(), 1)\n    fnew = f(xnew)\n    return np.concatenate([xnew[:, None], fnew[:, None]], axis = -1).astype(int)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# data = train_annotations[\"data\"][0]\n# data = np.array(ast.literal_eval(data))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# interpolate_mask(data).shape","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Classes are reduced down to ETT, NGT, CVC, Swan Ganz. With some slight modification this model could do each in their own class and then try to use it to directly classify the catheter/lines, but for this usage it is likely reasonable to just group them together to try to get the most accurate catheter tracing first. "},{"metadata":{"execution":{"iopub.execute_input":"2021-01-05T05:34:23.964149Z","iopub.status.busy":"2021-01-05T05:34:23.961043Z","iopub.status.idle":"2021-01-05T05:34:23.966851Z","shell.execute_reply":"2021-01-05T05:34:23.966359Z"},"papermill":{"duration":0.024921,"end_time":"2021-01-05T05:34:23.966948","exception":false,"start_time":"2021-01-05T05:34:23.942027","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"COLOR_MAP = {'ETT - Abnormal': 0,\n             'ETT - Borderline': 0,\n             'ETT - Normal': 0,\n             'NGT - Abnormal': 1,\n             'NGT - Borderline': 1,\n             'NGT - Incompletely Imaged': 1,\n             'NGT - Normal': 1,\n             'CVC - Abnormal': 2,\n             'CVC - Borderline': 2,\n             'CVC - Normal': 2,\n             'Swan Ganz Catheter Present': 3,\n            }\n\n\nclass SegDataset(Dataset):\n    def __init__(self, df, df_annotations, annot_size=10, transform=None, mode = 'train'):\n        self.df = df\n        self.df_annotations = df_annotations\n        self.annot_size = annot_size\n        self.file_names = df['file_path'].values\n        self.patient_id = df['StudyInstanceUID'].values\n        self.labels = df[target_cols].values\n        self.transform = transform\n        self.mode = mode\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        file_name = self.file_names[idx]\n        patient_id = self.patient_id[idx]\n        no_anno = 1\n        image = cv2.imread(file_name)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        seg_mask = np.zeros((image.shape[0], image.shape[1], 4), dtype = np.float32)\n        query_string = f\"StudyInstanceUID == '{patient_id}'\"\n        df = self.df_annotations.query(query_string)\n        if len(df) == 0:\n            no_anno = 0\n        for i, row in df.iterrows():\n            label = row[\"label\"]\n            data = np.array(ast.literal_eval(row[\"data\"]))\n            for data_point in range(len(data)):\n                point_pairs = data[data_point: data_point + 2]\n                if len(point_pairs) < 2:\n                    continue\n                for d in interpolate_mask(point_pairs):\n                    seg_mask[d[1]-self.annot_size//2:d[1]+self.annot_size//2,\n                          d[0]-self.annot_size//2:d[0]+self.annot_size//2,\n                          COLOR_MAP[label]] = 1\n        if self.transform:\n            augmented = self.transform(image=image, mask = seg_mask)\n            image = augmented['image']\n            mask = augmented['mask']\n        if self.mode == 'test':\n            return torch.tensor(image).float()\n        else:\n            label = torch.tensor(self.labels[idx]).float()\n            no_anno = torch.tensor(no_anno)\n            return file_name, torch.tensor(image).float(), torch.tensor(mask).float(), label, no_anno.float()","execution_count":null,"outputs":[]},{"metadata":{"execution":{"iopub.execute_input":"2021-01-05T05:34:23.999717Z","iopub.status.busy":"2021-01-05T05:34:23.999204Z","iopub.status.idle":"2021-01-05T05:34:24.164253Z","shell.execute_reply":"2021-01-05T05:34:24.163489Z"},"papermill":{"duration":0.185474,"end_time":"2021-01-05T05:34:24.164369","exception":false,"start_time":"2021-01-05T05:34:23.978895","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"train_annotations = pd.read_csv('../input/ranzcr-clip-catheter-line-classification/train_annotations.csv')\ndf_train = pd.read_csv('../input/how-to-properly-split-folds/train_folds.csv')\ndf_train['file_path'] = df_train.StudyInstanceUID.apply(lambda x: os.path.join(data_dir, f'{x}.jpg'))\nif debug:\n    df_train = df_train.sample(frac=0.1)\ntarget_cols = df_train.iloc[:, 1:12].columns.tolist()\n","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.011774,"end_time":"2021-01-05T05:34:24.188182","exception":false,"start_time":"2021-01-05T05:34:24.176408","status":"completed"},"tags":[]},"cell_type":"markdown","source":"## Utils"},{"metadata":{},"cell_type":"markdown","source":"Training is applied only to the samples that we have annotations for, but validation is done on all validation fold samples. Only samples with annotations are scored against but all are predicted for and then written to .npy files for later inspection and reuse. "},{"metadata":{"execution":{"iopub.execute_input":"2021-01-05T05:34:24.229730Z","iopub.status.busy":"2021-01-05T05:34:24.229182Z","iopub.status.idle":"2021-01-05T05:34:24.232895Z","shell.execute_reply":"2021-01-05T05:34:24.232419Z"},"papermill":{"duration":0.033092,"end_time":"2021-01-05T05:34:24.233024","exception":false,"start_time":"2021-01-05T05:34:24.199932","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"def macro_multilabel_auc(label, pred):\n    aucs = []\n    for i in range(len(target_cols)):\n        aucs.append(roc_auc_score(label[:, i], pred[:, i]))\n    print(np.round(aucs, 4))\n    return np.mean(aucs)\n\n\ndef train_func(train_loader):\n    model.train()\n    bar = tqdm(train_loader)\n    if use_amp:\n        scaler = torch.cuda.amp.GradScaler()\n    losses = []\n    seg_losses = []\n    for batch_idx, (_, images, masks, targets, no_annos) in enumerate(bar):\n        images, masks, targets, no_annos = images.to(device), masks.to(device), targets.to(device), no_annos.to(device)\n        with torch.cuda.amp.autocast():\n            mask_logits, logits = model(images)\n            loss = criterion(logits, targets).mean()\n            seg_loss = criterion(mask_logits, masks)\n            seg_loss = seg_loss.mean(axis = -1).mean(axis = -1).mean(axis = -1)\n            seg_loss = (seg_loss * no_annos)\n            seg_loss = seg_loss.mean()\n            total_loss = seg_loss #+ loss\n        scaler.scale(total_loss).backward()\n        scaler.step(optimizer)\n        scaler.update()\n        optimizer.zero_grad()\n        \n        losses.append(loss.item())\n        smooth_loss = np.mean(losses)\n        \n        seg_losses.append(seg_loss.item())\n        smooth_seg_loss = np.mean(seg_losses)\n        bar.set_description(f'loss: {loss.item():.5f}, smth: {smooth_loss:.5f}, seg_smth: {smooth_seg_loss:.5f}')\n\n    loss_train = np.mean(losses)\n    return loss_train\n\n\ndef valid_func(valid_loader):\n    model.eval()\n    bar = tqdm(valid_loader)\n\n    PROB = []\n    TARGETS = []\n    losses = []\n    seg_losses = []\n    PREDS = []\n    os.makedirs(f'./models/{model_name}/{fold_id}/seg_masks/', exist_ok = True)\n    with torch.no_grad():\n        for batch_idx, (filenames, images, masks, targets, no_annos) in enumerate(bar):\n            images, masks, targets, no_annos = images.to(device), masks.to(device), targets.to(device), no_annos.to(device)\n            mask_logits, logits = model(images)\n            for idx, filename in enumerate(filenames):\n                filename = \"\".join(filename.split(\"/\")[-1:])\n                filename = \".\".join(filename.split(\".\")[:-1])\n                mask_pred = mask_logits[idx].sigmoid().detach().cpu().numpy()*255\n                mask_pred = mask_pred.astype(np.uint8)\n                np.save(f'./models/{model_name}/{fold_id}/seg_masks/{filename}.npy',\n                        mask_pred)\n            PREDS += [logits.sigmoid()]\n            TARGETS += [targets.detach().cpu()]\n            loss = criterion(logits, targets).mean()\n            losses.append(loss.item())\n            smooth_loss = np.mean(losses)\n            \n            seg_loss = criterion(mask_logits, masks).mean(axis = -1).mean(axis = -1).mean(axis = -1)\n            if no_annos.sum() != 0:\n                seg_loss = (seg_loss * no_annos).sum()/no_annos.sum()\n                seg_losses.append(seg_loss.item())\n                seg_smooth_loss = np.mean(seg_losses)\n                bar.set_description(f'loss: {loss.item():.5f}, smth: {smooth_loss:.5f}, seg_smth: {seg_smooth_loss:.5f}')\n            \n    PREDS = torch.cat(PREDS).cpu().numpy()\n    TARGETS = torch.cat(TARGETS).cpu().numpy()\n    #roc_auc = roc_auc_score(TARGETS.reshape(-1), PREDS.reshape(-1))\n    roc_auc = macro_multilabel_auc(TARGETS, PREDS)\n    loss_valid = np.mean(losses)\n    seg_loss = seg_smooth_loss\n    return loss_valid, roc_auc, seg_loss","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.011616,"end_time":"2021-01-05T05:34:24.256758","exception":false,"start_time":"2021-01-05T05:34:24.245142","status":"completed"},"tags":[]},"cell_type":"markdown","source":"## Training"},{"metadata":{"execution":{"iopub.execute_input":"2021-01-05T05:34:24.284780Z","iopub.status.busy":"2021-01-05T05:34:24.284189Z","iopub.status.idle":"2021-01-05T05:34:33.404981Z","shell.execute_reply":"2021-01-05T05:34:33.403871Z"},"papermill":{"duration":9.136543,"end_time":"2021-01-05T05:34:33.405097","exception":false,"start_time":"2021-01-05T05:34:24.268554","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"def make_model(model_name):\n    model = SegNet(out_dim=len(target_cols),model_name = model_name, pretrained=True)\n    model = torch.nn.DataParallel(model)\n    model = model.to(device)\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"init_lr = 1e-3\nbatch_size = 8\nvalid_batch_size = 8\nn_epochs = 1\nnum_workers = 4\nearly_stop = 5","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def save_state(fp, model, scheduler, optimizer, epoch):\n    model.eval()\n    torch.save({\n                \"model\":model.state_dict(),\n                \"scheduler\":scheduler.state_dict(),\n                \"optimizer\":optimizer.state_dict(),\n                \"epoch\": epoch\n               },\n                fp\n              )\ndef load_state(fp, model, scheduler, optimizer):\n    state = torch.load(fp)\n    model.load_state_dict(state[\"model\"])\n    scheduler.load_state_dict(state[\"scheduler\"])\n    optimizer.load_state_dict(state[\"optimizer\"])\n    epoch = state[\"epoch\"]\n    return model, scheduler, optimizer, epoch","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# for item in dataset_train:\n#     break","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# plt.imshow(item[1][:, :, 10].detach().cpu())","execution_count":null,"outputs":[]},{"metadata":{"execution":{"iopub.execute_input":"2021-01-05T05:34:33.586772Z","iopub.status.busy":"2021-01-05T05:34:33.586238Z","iopub.status.idle":"2021-01-05T05:40:16.746117Z","shell.execute_reply":"2021-01-05T05:40:16.745539Z"},"papermill":{"duration":343.185712,"end_time":"2021-01-05T05:40:16.746245","exception":false,"start_time":"2021-01-05T05:34:33.560533","status":"completed"},"scrolled":true,"tags":[],"trusted":true},"cell_type":"code","source":"for fold_id in range(5)[:1]:\n    log = {}\n    roc_auc_max = 0.\n    seg_loss_min = 100\n    loss_min = 99999\n    not_improving = 0\n    df_train_this = df_train[df_train['fold'] != fold_id]\n    df_valid_this = df_train[df_train['fold'] == fold_id]\n\n    df_train_this = df_train_this[df_train_this['StudyInstanceUID'].isin(train_annotations['StudyInstanceUID'].unique())].reset_index(drop=True)\n\n    dataset_train = SegDataset(df_train_this, train_annotations, annot_size = 20, transform=transforms_train, mode = 'train')\n    dataset_valid = SegDataset(df_valid_this, train_annotations, annot_size = 20, transform=transforms_valid, mode = 'train')\n\n    train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=batch_size, shuffle=True,  num_workers=num_workers, pin_memory=True)\n    valid_loader = torch.utils.data.DataLoader(dataset_valid, batch_size=valid_batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)\n    \n    model = make_model(model_name)\n    criterion = nn.BCEWithLogitsLoss()\n    optimizer = optim.Adam(model.parameters(), lr=init_lr)\n    \n    scheduler_cosine = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, n_epochs, eta_min=1e-7)\n    model_path = f'./models/{model_name}/{fold_id}/{model_name}_fold{fold_id}_{phase}_best_loss.bin'\n    if os.path.exists(model_path):\n        print(\"model is resuming training\")\n        model, scheduler_cosine, optimizer, epoch = load_state(model_path, model, scheduler_cosine, optimizer)\n    else:\n        print(\"model being trained from scratch\")\n        epoch = 1\n    for epoch in range(epoch, n_epochs+1):\n        scheduler_cosine.step(epoch-1)\n        loss_train = train_func(train_loader)\n        loss_valid, roc_auc, seg_loss = valid_func(valid_loader)\n\n        log['loss_train'] = log.get('loss_train', []) + [loss_train]\n        log['loss_valid'] = log.get('loss_valid', []) + [loss_valid]\n        log['lr'] = log.get('lr', []) + [optimizer.param_groups[0][\"lr\"]]\n        log['roc_auc'] = log.get('roc_auc', []) + [roc_auc]\n        log['seg_loss'] = log.get('seg_loss', []) + [seg_loss]\n\n        content = time.ctime() + ' ' + f'Fold {fold_id}, Epoch {epoch}, lr: {optimizer.param_groups[0][\"lr\"]:.7f}, loss_train: {loss_train:.5f}, loss_valid: {loss_valid:.5f}, roc_auc: {roc_auc:.6f}.'\n        print(content)\n        not_improving += 1\n        os.makedirs(f'./models/{model_name}/{fold_id}/', exist_ok = True)\n        if seg_loss < seg_loss_min:\n            print(f'seg_loss_min ({seg_loss_min:.6f} --> {seg_loss:.6f}). Saving model ...')\n            save_state(f'./models/{model_name}/{fold_id}/{model_name}_fold{fold_id}_{phase}_best_AUC.bin',\n                       model,\n                       scheduler_cosine,\n                       optimizer,\n                       epoch)\n            seg_loss_min = seg_loss\n            not_improving = 0\n\n        if loss_valid < loss_min:\n            loss_min = loss_valid\n            save_state(f'./models/{model_name}/{fold_id}/{model_name}_fold{fold_id}_{phase}_best_loss.bin',\n                       model,\n                       scheduler_cosine,\n                       optimizer,\n                       epoch)\n        if not_improving == early_stop:\n            print('Early Stopping...')\n            break\n    log_df = pd.DataFrame(log)\n    log_df.to_csv(f'./models/{model_name}/{fold_id}/logs_{fold_id}_{phase}.csv')\n    df_valid_this.to_csv(f'./models/{model_name}/{fold_id}/val_df_{fold_id}_{phase}.csv')\n    \n    save_state(f'./models/{model_name}/{fold_id}/{model_name}_fold{fold_id}_{phase}_final.bin',\n           model,\n           scheduler_cosine,\n           optimizer,\n           epoch)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"img = np.load(\"./models/resnet18/0/seg_masks/1.2.826.0.1.3680043.8.498.11073617724281949099281046870716891732.npy\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"img.shape","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Visualization of the models predictions for the 4 different classes. Results are likely not great since not fully trained. You can try training for more epochs and with larger models than resnet18 to get better looking masks. "},{"metadata":{"trusted":true},"cell_type":"code","source":"for i in range(4):\n    plt.imshow(img[:, :, i])\n    plt.show()\n    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","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}