{"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":"import torch\nfrom PIL import Image\nfrom sklearn.model_selection import train_test_split\nfrom torch import nn\nfrom torch import optim\nimport torch.nn.functional as F\nimport torchvision.transforms.functional as transF\nimport numpy as np\nfrom torch.utils.data import DataLoader, Dataset\nfrom torchvision import datasets, transforms, models\nfrom torchvision.datasets import ImageFolder\nimport matplotlib.pyplot as plt\nimport matplotlib.ticker as mticker\nimport os\nfrom albumentations import (HorizontalFlip, ShiftScaleRotate, Normalize, Resize, Compose, GaussNoise)\nimport cv2\nfrom albumentations.pytorch import ToTensorV2\nfrom skimage import io, transform\nfrom skimage.io import imread, imshow\nimport cv2\nfrom PIL import Image\nfrom torch.utils.data import random_split\nfrom tqdm import tqdm","metadata":{"execution":{"iopub.status.busy":"2022-02-18T02:45:42.877239Z","iopub.execute_input":"2022-02-18T02:45:42.878063Z","iopub.status.idle":"2022-02-18T02:45:46.534909Z","shell.execute_reply.started":"2022-02-18T02:45:42.877931Z","shell.execute_reply":"2022-02-18T02:45:46.534165Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!mkdir train_data test_data\n!unzip -qq ../input/data-science-bowl-2018/stage1_train.zip -d train_data\n!unzip -qq ../input/data-science-bowl-2018/stage1_test.zip -d test_data","metadata":{"execution":{"iopub.status.busy":"2022-02-18T02:45:46.536750Z","iopub.execute_input":"2022-02-18T02:45:46.536987Z","iopub.status.idle":"2022-02-18T02:45:51.194171Z","shell.execute_reply.started":"2022-02-18T02:45:46.536953Z","shell.execute_reply":"2022-02-18T02:45:51.193090Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(device)","metadata":{"execution":{"iopub.status.busy":"2022-02-18T02:45:51.196223Z","iopub.execute_input":"2022-02-18T02:45:51.196814Z","iopub.status.idle":"2022-02-18T02:45:51.245762Z","shell.execute_reply.started":"2022-02-18T02:45:51.196767Z","shell.execute_reply":"2022-02-18T02:45:51.244920Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CFG:\n    size = 128\n    mean = [43.5334, 39.5612, 48.2250]\n    std = [67.0821, 61.8689, 75.3638]\n    batch_size = 10","metadata":{"execution":{"iopub.status.busy":"2022-02-18T02:45:51.249131Z","iopub.execute_input":"2022-02-18T02:45:51.249418Z","iopub.status.idle":"2022-02-18T02:45:51.254731Z","shell.execute_reply.started":"2022-02-18T02:45:51.249382Z","shell.execute_reply":"2022-02-18T02:45:51.253886Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_transforms(mean=CFG.mean, std=CFG.std):\n            return Compose([\n                            HorizontalFlip(p=0.5),\n                            Normalize(mean=mean, std=std, p=1),\n                            ToTensorV2(transpose_mask=True)\n                            ])","metadata":{"execution":{"iopub.status.busy":"2022-02-18T02:45:51.259052Z","iopub.execute_input":"2022-02-18T02:45:51.259289Z","iopub.status.idle":"2022-02-18T02:45:51.265238Z","shell.execute_reply.started":"2022-02-18T02:45:51.259256Z","shell.execute_reply":"2022-02-18T02:45:51.264292Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CustomDataset(Dataset):\n    def __init__(self,path):\n        self.path = path\n        self.folders = os.listdir(path)\n        print(len(self.folders))\n        self.transforms = get_transforms()\n    def __getitem__(self,idx):\n        image_folder = os.path.join(self.path,self.folders[idx],'images/')\n        mask_folder = os.path.join(self.path,self.folders[idx],'masks/')\n        image_path = os.path.join(image_folder,os.listdir(image_folder)[0])\n    \n        img = io.imread(image_path)[:,:,:3].astype('float32')\n        img = transform.resize(img,(128,128))\n        mask = self.get_mask(mask_folder, 128 ).astype('float32')\n    \n        augmented = self.transforms(image=img, mask=mask)\n        img = augmented['image']\n        mask = augmented['mask']\n        return (img,mask) \n\n    def get_mask(self,mask_folder, img_size):\n        mask = np.zeros((img_size, img_size, 1), dtype=np.bool)\n        for mask_ in os.listdir(mask_folder):\n            mask_ = io.imread(os.path.join(mask_folder,mask_))\n            mask_ = transform.resize(mask_, (img_size, img_size))\n            mask_ = np.expand_dims(mask_,axis=-1)\n            mask = np.maximum(mask, mask_)\n        return mask\n\n    def __len__(self):\n        return len(self.folders)","metadata":{"execution":{"iopub.status.busy":"2022-02-18T02:45:51.266714Z","iopub.execute_input":"2022-02-18T02:45:51.267313Z","iopub.status.idle":"2022-02-18T02:45:51.280558Z","shell.execute_reply.started":"2022-02-18T02:45:51.267252Z","shell.execute_reply":"2022-02-18T02:45:51.279762Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"base_dir = './train_data'\ndata = CustomDataset(base_dir)\nfor image,mask in data:\n    print(image.shape)\n    print(mask.shape)\n    break","metadata":{"execution":{"iopub.status.busy":"2022-02-18T04:36:13.594747Z","iopub.execute_input":"2022-02-18T04:36:13.595036Z","iopub.status.idle":"2022-02-18T04:36:13.738928Z","shell.execute_reply.started":"2022-02-18T04:36:13.594991Z","shell.execute_reply":"2022-02-18T04:36:13.738105Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class UnNormalize(object):\n    def __init__(self, mean, std):\n        self.mean = mean\n        self.std = std\n\n    def __call__(self, tensor):\n        for t, m, s in zip(tensor, self.mean, self.std):\n            t.mul_(s).add_(m)\n        return tensor\n\nunorm = UnNormalize(mean = CFG.mean, std = CFG.std)","metadata":{"execution":{"iopub.status.busy":"2022-02-18T02:45:51.523760Z","iopub.execute_input":"2022-02-18T02:45:51.524167Z","iopub.status.idle":"2022-02-18T02:45:51.534409Z","shell.execute_reply.started":"2022-02-18T02:45:51.524126Z","shell.execute_reply":"2022-02-18T02:45:51.533279Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_len = int(0.9*len(data))\nval_len = len(data)-train_len\ntrain_data, val_data = random_split(data, [train_len,val_len])\ntrain_loader = DataLoader(dataset=train_data, batch_size=CFG.batch_size, shuffle=True)\nval_loader = DataLoader(dataset=val_data, batch_size=CFG.batch_size)\nprint(len(train_loader),len(val_loader))","metadata":{"execution":{"iopub.status.busy":"2022-02-18T02:45:51.535789Z","iopub.execute_input":"2022-02-18T02:45:51.536400Z","iopub.status.idle":"2022-02-18T02:45:51.556118Z","shell.execute_reply.started":"2022-02-18T02:45:51.536361Z","shell.execute_reply":"2022-02-18T02:45:51.555426Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.figure(figsize=(20,20))\nfor (images,labels) in train_loader:\n    for i, (image, label) in enumerate(zip(images,labels)):\n        plt.subplot(10, 2, 2*i+1)\n        plt.axis('off')\n        temp = unorm(image.clone())\n        plt.title('image')\n        plt.imshow( temp.permute(1, 2, 0))\n        \n        plt.subplot(10, 2, 2*i+2)\n        plt.axis('off')\n        temp = unorm(label.clone())\n        plt.title('mask')\n        plt.imshow( temp.squeeze(0),cmap='gray')\n    plt.show()\n    break","metadata":{"execution":{"iopub.status.busy":"2022-02-18T02:45:51.557863Z","iopub.execute_input":"2022-02-18T02:45:51.560157Z","iopub.status.idle":"2022-02-18T02:45:55.628402Z","shell.execute_reply.started":"2022-02-18T02:45:51.560114Z","shell.execute_reply":"2022-02-18T02:45:55.626988Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def calculate_iou(y_pred,y):\n    inputs = y_pred.reshape(-1)\n    targets = y.reshape(-1)\n    intersection = (inputs * targets).sum()\n    total = (inputs + targets).sum()\n    union = total - intersection \n    smooth = 1    \n    iou = (intersection + smooth)/(union + smooth)\n    return iou\n\n\n\ndef calculate_iou_batch(y_pred,y):\n    ious = []\n    y_pred = torch.sigmoid(y_pred)\n    y_pred = y_pred.clone().cpu().detach().numpy()\n    y = y.clone().cpu().detach().numpy() \n    \n    for pred, label in zip(y_pred, y):\n        ious.append(calculate_iou(pred, label))\n    iou = np.nanmean(ious)\n    return iou   ","metadata":{"execution":{"iopub.status.busy":"2022-02-18T02:45:55.629422Z","iopub.execute_input":"2022-02-18T02:45:55.629652Z","iopub.status.idle":"2022-02-18T02:45:55.641349Z","shell.execute_reply.started":"2022-02-18T02:45:55.629622Z","shell.execute_reply":"2022-02-18T02:45:55.640680Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class DoubleConv(nn.Module):\n    def __init__(self, in_channels, out_channels):\n        super(DoubleConv, self).__init__()\n        self.double_conv = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(inplace=True),\n        )\n\n    def forward(self, x):\n        return self.double_conv(x)\n    \n    \nclass DownBlock(nn.Module):\n    def __init__(self, in_channels, out_channels):\n        super(DownBlock, self).__init__()\n        self.double_conv = DoubleConv(in_channels, out_channels)\n        self.down_sample = nn.MaxPool2d(2)\n\n    def forward(self, x):\n        skip_out = self.double_conv(x)\n        down_out = self.down_sample(skip_out)\n        return (down_out, skip_out)\n\n    \nclass UpBlock(nn.Module):\n    def __init__(self, in_channels, out_channels, up_sample_mode):\n        super(UpBlock, self).__init__()\n        if up_sample_mode == 'conv_transpose':\n            self.up_sample = nn.ConvTranspose2d(in_channels-out_channels, in_channels-out_channels, kernel_size=2, stride=2)        \n        elif up_sample_mode == 'bilinear':\n            self.up_sample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)\n        else:\n            raise ValueError(\"Unsupported `up_sample_mode` (can take one of `conv_transpose` or `bilinear`)\")\n        self.double_conv = DoubleConv(in_channels, out_channels)\n\n    def forward(self, down_input, skip_input):\n        x = self.up_sample(down_input)\n        x = torch.cat([x, skip_input], dim=1)\n        return self.double_conv(x)\n\n    \nclass UNet(nn.Module):\n    def __init__(self, out_classes=1, up_sample_mode='conv_transpose'):\n        super(UNet, self).__init__()\n        self.up_sample_mode = up_sample_mode\n        # Downsampling Path\n        self.down_conv1 = DownBlock(3, 64)\n        self.down_conv2 = DownBlock(64, 128)\n        self.down_conv3 = DownBlock(128, 256)\n        self.down_conv4 = DownBlock(256, 512)\n        # Bottleneck\n        self.double_conv = DoubleConv(512, 1024)\n        # Upsampling Path\n        self.up_conv4 = UpBlock(512 + 1024, 512, self.up_sample_mode)\n        self.up_conv3 = UpBlock(256 + 512, 256, self.up_sample_mode)\n        self.up_conv2 = UpBlock(128 + 256, 128, self.up_sample_mode)\n        self.up_conv1 = UpBlock(128 + 64, 64, self.up_sample_mode)\n        # Final Convolution\n        self.conv_last = nn.Conv2d(64, out_classes, kernel_size=1)\n\n    def forward(self, x):\n        x, skip1_out = self.down_conv1(x)\n        x, skip2_out = self.down_conv2(x)\n        x, skip3_out = self.down_conv3(x)\n        x, skip4_out = self.down_conv4(x)\n        x = self.double_conv(x)\n        x = self.up_conv4(x, skip4_out)\n        x = self.up_conv3(x, skip3_out)\n        x = self.up_conv2(x, skip2_out)\n        x = self.up_conv1(x, skip1_out)\n        x = self.conv_last(x)\n        return x","metadata":{"execution":{"iopub.status.busy":"2022-02-18T02:45:55.643077Z","iopub.execute_input":"2022-02-18T02:45:55.643694Z","iopub.status.idle":"2022-02-18T02:45:55.662459Z","shell.execute_reply.started":"2022-02-18T02:45:55.643632Z","shell.execute_reply":"2022-02-18T02:45:55.661758Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = UNet().to(device)\nimage = torch.rand((1,3,128,128)).to(device)\nmodel(image).shape","metadata":{"execution":{"iopub.status.busy":"2022-02-18T02:45:55.663727Z","iopub.execute_input":"2022-02-18T02:45:55.664361Z","iopub.status.idle":"2022-02-18T02:46:04.035719Z","shell.execute_reply.started":"2022-02-18T02:45:55.664334Z","shell.execute_reply":"2022-02-18T02:46:04.034946Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class DiceBCELoss(nn.Module):\n    def __init__(self, weight=None, size_average=True):\n        super(DiceBCELoss, self).__init__()\n\n    def forward(self, inputs, targets, smooth=1):\n        \n        #comment out if your model contains a sigmoid or equivalent activation layer\n        inputs = torch.sigmoid(inputs)       \n        bce_weight = 0.5\n        #flatten label and prediction tensors\n        inputs = inputs.view(-1)\n        targets = targets.view(-1)\n        \n        intersection = (inputs * targets).sum()                            \n        dice_loss = 1 - (2.*intersection + smooth)/(inputs.sum() + targets.sum() + smooth)  \n        BCE = F.binary_cross_entropy(inputs, targets, reduction='mean')\n        loss_final = BCE * bce_weight + dice_loss * (1 - bce_weight)\n        return loss_final ","metadata":{"execution":{"iopub.status.busy":"2022-02-18T02:46:04.039347Z","iopub.execute_input":"2022-02-18T02:46:04.039570Z","iopub.status.idle":"2022-02-18T02:46:04.047463Z","shell.execute_reply.started":"2022-02-18T02:46:04.039542Z","shell.execute_reply":"2022-02-18T02:46:04.046526Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"learning_rate = 3e-4\noptimizer = optim.Adam(model.parameters(),lr= learning_rate)\ncriterion = DiceBCELoss()\nepochs = 25","metadata":{"execution":{"iopub.status.busy":"2022-02-18T02:46:04.048778Z","iopub.execute_input":"2022-02-18T02:46:04.049578Z","iopub.status.idle":"2022-02-18T02:46:04.057815Z","shell.execute_reply.started":"2022-02-18T02:46:04.049539Z","shell.execute_reply":"2022-02-18T02:46:04.057078Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train(model, train_loader):\n    model.train()\n    train_loss, train_score, target_count = 0, 0, 0\n    \n    for images, masks in tqdm(train_loader):\n        images = images.to(device)\n        masks  = masks.to(device)\n\n        # Forward pass\n        outputs = model(images)\n        loss = criterion(outputs, masks)\n        score = calculate_iou_batch(outputs,masks)\n        train_loss += loss.item()\n        train_score += score\n        target_count += masks.shape[0]\n        # Backward and optimize\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n    return train_loss/target_count, train_score / target_count","metadata":{"execution":{"iopub.status.busy":"2022-02-18T02:46:04.058800Z","iopub.execute_input":"2022-02-18T02:46:04.061139Z","iopub.status.idle":"2022-02-18T02:46:04.068873Z","shell.execute_reply.started":"2022-02-18T02:46:04.061097Z","shell.execute_reply":"2022-02-18T02:46:04.067913Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def validate(model,val_loader):\n    model.eval()\n    val_loss, val_score, target_count = 0, 0, 0\n    for images, masks in val_loader:\n        images = images.to(device)\n        masks = masks.to(device)\n\n        outputs = model(images)\n        loss = criterion(outputs, masks)\n        score = calculate_iou_batch(outputs,masks)\n        val_loss += loss.item()\n        val_score += score\n        target_count += masks.shape[0]\n\n    return val_loss/target_count, val_score / target_count","metadata":{"execution":{"iopub.status.busy":"2022-02-18T02:46:04.070409Z","iopub.execute_input":"2022-02-18T02:46:04.071316Z","iopub.status.idle":"2022-02-18T02:46:04.082374Z","shell.execute_reply.started":"2022-02-18T02:46:04.071273Z","shell.execute_reply":"2022-02-18T02:46:04.081567Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_loss_history,val_loss_history = [],[]\ntrain_iou_history,val_iou_history = [],[]\nmin_val_loss = 1e9\ncheckpoint_path = 'checkpoint'\n\nfor epoch in range(epochs):\n    train_loss, train_score = train(model, train_loader)\n    val_loss, val_score = validate(model, val_loader)\n    train_loss_history.append(train_loss)\n    val_loss_history.append(val_loss)\n    train_iou_history.append(train_score)\n    val_iou_history.append(val_score)\n    print(\"Epoch {0}: train_score {1} \\t train_loss {2} \\t val_score {3} \\t val_loss {4}\".format(epoch, train_score, train_loss, val_score, val_loss))\n    if val_loss < min_val_loss:\n        checkpoint = {\n                'epoch': epoch + 1,\n                'valid_loss_min': val_loss,\n                'state_dict': model.state_dict(),\n                'optimizer': optimizer.state_dict(),\n            }\n        torch.save(checkpoint,checkpoint_path)\n        min_val_loss = val_loss","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2022-02-18T02:46:04.083861Z","iopub.execute_input":"2022-02-18T02:46:04.084531Z","iopub.status.idle":"2022-02-18T04:26:52.266166Z","shell.execute_reply.started":"2022-02-18T02:46:04.084490Z","shell.execute_reply":"2022-02-18T04:26:52.265391Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.figure(figsize=(25,5))\nplt.plot(train_loss_history,'-o')\nplt.plot(val_loss_history,'-o')\nplt.xlabel('epochs')\nplt.ylabel('losses')\nplt.gca().xaxis.set_major_locator(mticker.MultipleLocator(1))\nplt.legend(['Train Loss','Val Loss'])\nplt.title('Train Loss and Val Loss')\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-02-18T04:26:52.267590Z","iopub.execute_input":"2022-02-18T04:26:52.268034Z","iopub.status.idle":"2022-02-18T04:26:52.542474Z","shell.execute_reply.started":"2022-02-18T04:26:52.267983Z","shell.execute_reply":"2022-02-18T04:26:52.541833Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.figure(figsize=(25,5))\nplt.plot(train_iou_history,'-o')\nplt.plot(val_iou_history,'-o')\nplt.xlabel('epochs')\nplt.ylabel('scores')\nplt.gca().xaxis.set_major_locator(mticker.MultipleLocator(1))\nplt.legend(['Train Score','Val Score'])\nplt.title('Train Score and Val Score')\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-02-18T04:26:52.543913Z","iopub.execute_input":"2022-02-18T04:26:52.544397Z","iopub.status.idle":"2022-02-18T04:26:52.822585Z","shell.execute_reply.started":"2022-02-18T04:26:52.544355Z","shell.execute_reply":"2022-02-18T04:26:52.821848Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"checkpoint = torch.load(checkpoint_path)\nmodel.load_state_dict(checkpoint['state_dict'])","metadata":{"execution":{"iopub.status.busy":"2022-02-18T04:26:52.823698Z","iopub.execute_input":"2022-02-18T04:26:52.824014Z","iopub.status.idle":"2022-02-18T04:26:53.090696Z","shell.execute_reply.started":"2022-02-18T04:26:52.823977Z","shell.execute_reply":"2022-02-18T04:26:53.089912Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"images,masks = next(iter(val_loader))\nprint(images.shape,masks.shape)","metadata":{"execution":{"iopub.status.busy":"2022-02-18T04:26:53.094607Z","iopub.execute_input":"2022-02-18T04:26:53.096612Z","iopub.status.idle":"2022-02-18T04:26:56.722525Z","shell.execute_reply.started":"2022-02-18T04:26:53.096570Z","shell.execute_reply":"2022-02-18T04:26:56.721710Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"images = images.to(device)\npredict_masks = model(images)\nprint(images.shape, masks.shape, predict_masks.shape)\nplt.figure(figsize=(20,20))\nfor i, (image, mask, predict_mask) in enumerate(zip(images,masks,predict_masks)):\n    plt.subplot(10, 3, i*3+1)\n    plt.axis('off')\n    temp = unorm(image.clone()).cpu()\n    plt.title('image')\n    plt.imshow( temp.permute(1, 2, 0))\n\n    plt.subplot(10, 3, i*3+2)\n    plt.axis('off')\n    temp = unorm(mask.clone()).cpu()\n    plt.title('Actual')\n    plt.imshow( temp.squeeze(0),cmap='gray')\n\n    plt.subplot(10, 3, i*3+3)\n    plt.axis('off')\n    temp = predict_mask.clone().cpu()\n    plt.title('Predict')\n    plt.imshow( temp.squeeze(0).detach().numpy(),cmap='gray')\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-02-18T04:26:56.723850Z","iopub.execute_input":"2022-02-18T04:26:56.724128Z","iopub.status.idle":"2022-02-18T04:26:58.080495Z","shell.execute_reply.started":"2022-02-18T04:26:56.724092Z","shell.execute_reply":"2022-02-18T04:26:58.078635Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class TestDataset(Dataset):\n    def __init__(self,path):\n        self.path = path\n        self.folders = os.listdir(path)\n        print(len(self.folders))\n        self.transforms = get_transforms()\n    def __getitem__(self,idx):\n        image_folder = os.path.join(self.path,self.folders[idx],'images/')\n        image_path = os.path.join(image_folder,os.listdir(image_folder)[0])\n    \n        img = io.imread(image_path)[:,:,:3].astype('float32')\n        img = transform.resize(img,(128,128))\n    \n        augmented = self.transforms(image=img)\n        img = augmented['image']\n        return img\n    def __len__(self):\n        return len(self.folders)","metadata":{"execution":{"iopub.status.busy":"2022-02-18T04:41:40.599953Z","iopub.execute_input":"2022-02-18T04:41:40.600609Z","iopub.status.idle":"2022-02-18T04:41:40.608366Z","shell.execute_reply.started":"2022-02-18T04:41:40.600569Z","shell.execute_reply":"2022-02-18T04:41:40.607544Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_dir = './test_data'\ntest_data = TestDataset(test_dir)\nfor image in test_data:\n    print(image.shape)\n    break","metadata":{"execution":{"iopub.status.busy":"2022-02-18T04:42:16.156772Z","iopub.execute_input":"2022-02-18T04:42:16.157582Z","iopub.status.idle":"2022-02-18T04:42:16.186142Z","shell.execute_reply.started":"2022-02-18T04:42:16.157535Z","shell.execute_reply":"2022-02-18T04:42:16.185354Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_loader = DataLoader(dataset=test_data, batch_size=CFG.batch_size)","metadata":{"execution":{"iopub.status.busy":"2022-02-18T04:43:04.701994Z","iopub.execute_input":"2022-02-18T04:43:04.702636Z","iopub.status.idle":"2022-02-18T04:43:04.707494Z","shell.execute_reply.started":"2022-02-18T04:43:04.702601Z","shell.execute_reply":"2022-02-18T04:43:04.706686Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"images = next(iter(test_loader))\nimages = images.to(device)\npredict_masks = model(images)\nprint(images.shape, predict_masks.shape)\nplt.figure(figsize=(20,20))\nfor i, (image, predict_mask) in enumerate(zip(images,predict_masks)):\n    plt.subplot(10, 2, i*2+1)\n    plt.axis('off')\n    temp = unorm(image.clone()).cpu()\n    plt.title('image')\n    plt.imshow( temp.permute(1, 2, 0))\n\n    plt.subplot(10, 2, i*2+2)\n    plt.axis('off')\n    temp = predict_mask.clone().cpu()\n    plt.title('Predict')\n    plt.imshow( temp.squeeze(0).detach().numpy(),cmap='gray')\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-02-18T04:47:53.847855Z","iopub.execute_input":"2022-02-18T04:47:53.848546Z","iopub.status.idle":"2022-02-18T04:47:55.067549Z","shell.execute_reply.started":"2022-02-18T04:47:53.848509Z","shell.execute_reply":"2022-02-18T04:47:55.066901Z"},"trusted":true},"execution_count":null,"outputs":[]}]}