{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"%reload_ext autoreload\n%autoreload 2\n%matplotlib inline","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","trusted":true},"cell_type":"code","source":"import os\nimport cv2\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport torch\nimport torchvision\nimport torchvision.transforms as transforms\nimport torchvision.datasets as datasets\nimport torch.optim as optim\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision.transforms.functional as TF\nfrom torch.utils.data.dataset import Dataset as Dataset\nfrom torch.utils.data import dataloader as DL\nfrom torch.autograd import Variable\nfrom PIL import Image\nfrom glob import glob\nimport PIL","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n\n# Assuming that we are on a CUDA machine, this should print a CUDA device:\n\nprint(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"input_folder = '../input'\n\ntrain_paths = glob('/'.join([input_folder,'train/*.jpg']))\ntrain_masks_paths = glob('/'.join([input_folder,'train_masks/*.gif']))\ntest_paths = glob('/'.join([input_folder,'test/*.jpg']))\n\ntrain_paths.sort()\ntrain_masks_paths.sort()\ntest_paths.sort()\n\ntemp1 = train_paths \ntrain_paths = train_paths[0:1]\ntemp2 = train_masks_paths\ntrain_masks_paths = train_masks_paths[0:1000]\n\ntest_paths = temp1[1100:1300]\ntest_masks_paths = temp2[1100:1300]\n\nprint('Number of training images: ', len(train_paths), 'Number of corresponding masks: ', len(train_masks_paths), 'Number of test images: ', len(test_paths))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class MyDataset(Dataset):\n    def __init__(self, image_paths, mask_paths, train=True):\n        self.image_paths = image_paths\n        self.mask_paths = mask_paths\n        \n    def transforms(self, image, mask):\n        #img = img.resize((wsize, baseheight), PIL.Image.ANTIALIAS)\n        #image = transforms.Resize(size=(64, 64))(image)\n        #mask = transforms.Resize(size=(64, 64))(mask)\n        image = image.resize((64, 64), PIL.Image.NEAREST)\n        mask = mask.resize((64, 64), PIL.Image.NEAREST)\n        image = TF.to_tensor(image)\n        mask = TF.to_tensor(mask)\n        \n        return image, mask\n    \n    def __getitem__(self, index):\n        image = Image.open(self.image_paths[index])\n        mask = Image.open(self.mask_paths[index])\n        \n        x, y = self.transforms(image, mask)\n        return x, y\n    \n    def __len__(self):\n        return len(self.image_paths)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(len(train_paths))\nprint(len(train_masks_paths))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dataset = MyDataset(train_paths, train_masks_paths)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"im, mask = train_dataset[0]\nim = im.numpy()\nmask = mask.numpy()\nprod = np.multiply(im, mask)\nprint(prod.shape)\nplt.imshow(prod[0])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(im)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(len(train_dataset))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_data_loader = DL.DataLoader(train_dataset, batch_size=1, shuffle=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(len(train_data_loader))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def double_conv(in_channels, out_channels):\n    return nn.Sequential(\n        nn.Conv2d(in_channels, out_channels, 3, padding=1),\n        nn.ReLU(inplace=True),\n        nn.Conv2d(out_channels, out_channels, 3, padding=1),\n        nn.ReLU(inplace=True)\n    )   \n\nclass UNet(nn.Module):\n\n    def __init__(self, n_class):\n        super().__init__()\n                \n        self.dconv_down1 = double_conv(3, 64)\n        self.dconv_down2 = double_conv(64, 128)\n        self.dconv_down3 = double_conv(128, 256)\n        self.dconv_down4 = double_conv(256, 512)  \n\n        self.maxpool = nn.MaxPool2d(2)\n        \n        self.dconv_up3 = double_conv(512, 256)\n        self.dconv_up2 = double_conv(256, 128)\n        self.dconv_up1 = double_conv(128, 64)\n        \n        self.TConv3 = nn.ConvTranspose2d(512, 256, 2, stride=2)\n        self.TConv2 = nn.ConvTranspose2d(256, 128, 2, stride=2)\n        self.TConv1 = nn.ConvTranspose2d(128, 64, 2, stride=2)\n        \n        self.conv_last = nn.Conv2d(64, n_class, 1)\n        \n    def forward(self, x):\n        conv1 = self.dconv_down1(x)\n        x = self.maxpool(conv1)\n        \n        conv2 = self.dconv_down2(x)\n        x = self.maxpool(conv2)\n        \n        conv3 = self.dconv_down3(x)\n        x = self.maxpool(conv3)\n        \n        x = self.dconv_down4(x)\n        \n        x = self.TConv3(x)\n        x = torch.cat([x, conv3], dim=1)\n\n        x = self.dconv_up3(x)\n        x = self.TConv2(x)\n        x = torch.cat([x, conv2], dim=1)\n\n        x = self.dconv_up2(x)\n        x = self.TConv1(x)\n        x = torch.cat([x, conv1], dim=1)\n\n        x = self.dconv_up1(x)\n        out = self.conv_last(x)\n        out = F.sigmoid(out)\n        \n        return out\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = UNet(n_class=1)\nmodel = model.cuda()\noptimizer = optim.Adam(model.parameters(), lr=3e-5)\ncriterion = nn.BCELoss()\nmodel.train()\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#Number of parameters in the model\nsum([p.numel() for p in model.parameters()])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dtype = torch.cuda.FloatTensor","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model.train()\n\nnum_epochs = 100\nrunning_loss = 0.0\n\nfor epoch in range(num_epochs):\n    running_loss = 0.0\n    for i, (X, y) in enumerate(train_data_loader):\n        X = X.to(device)\n        y = y.to(device)\n        X = Variable(X)\n        y = Variable(y)\n        \n        optimizer.zero_grad()\n        output = model(X)\n        loss = criterion(output, y)\n        \n        loss.backward()\n        optimizer.step()\n        running_loss += loss.item()\n    \n    print(\"loss for epoch \" + str(epoch) + \":  \" + str(running_loss))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Predict"},{"metadata":{"trusted":true},"cell_type":"code","source":"model.eval()\n\nx, y = train_dataset[0]\nx = x.unsqueeze(0)\nx = x.cuda()\npred = model(x)\n\npred = pred.squeeze(0)\npred = pred.squeeze(0)\n\nplt.imshow(pred.cpu().detach().numpy())\nprint(pred)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.6.4","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat":4,"nbformat_minor":1}