Rough Sketch Cleanup.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.autograd import Variable\n",
    "from torchvision import transforms\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from PIL import Image\n",
    "import random\n",
    "root=\"../own_dataset/split/\"\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "crop_size = 424\n",
    "batch_size = 8\n",
    "crop_num = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def default_loader(path):\n",
    "    return Image.open(path).convert(\"L\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyDataset(Dataset):\n",
    "    def __init__(self, txt, crop_size,crop_num, transform=None, target_transform=None, loader=default_loader):\n",
    "        fh = open(txt, 'r')\n",
    "        imgs = []\n",
    "        gts = []\n",
    "        lines = fh.read().split()\n",
    "        for i, line in enumerate(lines):\n",
    "            line = str(line)\n",
    "            for j in range(0,crop_num):\n",
    "                imgs.append('../own_dataset/input/'+line+'.jpg')\n",
    "                gts.append('../own_dataset/output/'+line+'.png')\n",
    "        self.imgs = imgs\n",
    "        self.gts = gts\n",
    "        self.transform = transform\n",
    "        self.target_transform = target_transform\n",
    "        self.loader = loader\n",
    "        self.crop = crop_size\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        fn = self.imgs[index]\n",
    "        img = self.loader(fn)\n",
    "        gt = self.gts[index]\n",
    "        gt_img = self.loader(gt)\n",
    "        w,h = gt_img.size\n",
    "        if w<crop_size or h<crop_size:\n",
    "            width = w\n",
    "            height = h\n",
    "            if w <= h:\n",
    "                width = crop_size\n",
    "                height = float(height/w) * crop_size\n",
    "                height = int(height)\n",
    "#                 print height\n",
    "            else:\n",
    "                height = crop_size\n",
    "                width = float(width / h) * crop_size\n",
    "                width = int(width)\n",
    "#                 print width\n",
    "            img = img.resize((width, height),Image.NEAREST)  \n",
    "            gt_img = gt_img.resize((width,height),Image.NEAREST)\n",
    "        w,h = gt_img.size\n",
    "        if self.crop !=-1:\n",
    "            x1, y1 = random.randint(0, w - self.crop), random.randint(0, h - self.crop)\n",
    "            img, gt_img = img.crop((x1, y1, x1 + self.crop, y1 + self.crop)), gt_img.crop((x1, y1, x1 + self.crop, y1 + self.crop))\n",
    "\n",
    "#         plt.figure()\n",
    "#         plt.imshow(gt_img,cmap ='gray')\n",
    "#         plt.figure()\n",
    "#         plt.imshow(img,cmap ='gray')\n",
    "        \n",
    "        if self.transform is not None:\n",
    "            img = self.transform(img)\n",
    "            gt_img = self.transform(gt_img)\n",
    "                  \n",
    "        return img,gt_img\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.imgs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data=MyDataset(root+'train2.txt', crop_size,crop_num, transform=transforms.ToTensor())\n",
    "train_loader = DataLoader(dataset=train_data, batch_size=batch_size, shuffle=True,drop_last = True)\n",
    "test_data=MyDataset(root+'val2.txt', -1, 1,transform=transforms.ToTensor())\n",
    "test_loader = DataLoader(dataset=test_data, batch_size=1)\n",
    "print(len(train_data))\n",
    "print(len(train_loader))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for batch_x, batch_y in train_loader:\n",
    "    print(batch_x.size())\n",
    "    print(batch_y.size())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        self.downconv1 = torch.nn.Sequential(\n",
    "            torch.nn.Conv2d(1, 48, 5, 2, 2),\n",
    "            torch.nn.BatchNorm2d(48),\n",
    "            torch.nn.ReLU(),\n",
    "            \n",
    "            torch.nn.Conv2d(48, 128, 3, 1, 1),\n",
    "            torch.nn.BatchNorm2d(128),\n",
    "            torch.nn.ReLU(),\n",
    "            \n",
    "            torch.nn.Conv2d(128, 128, 3, 1, 1),\n",
    "            torch.nn.BatchNorm2d(128),\n",
    "            torch.nn.ReLU(),\n",
    "        )\n",
    "        self.downconv2 = torch.nn.Sequential(\n",
    "            torch.nn.Conv2d(128,256, 3, 2, 1),\n",
    "            torch.nn.BatchNorm2d(256),\n",
    "            torch.nn.ReLU(),\n",
    "            \n",
    "            torch.nn.Conv2d(256, 256, 3, 1, 1),\n",
    "            torch.nn.BatchNorm2d(256),\n",
    "            torch.nn.ReLU(),\n",
    "            \n",
    "            torch.nn.Conv2d(256, 256, 3, 1, 1),\n",
    "            torch.nn.BatchNorm2d(256),\n",
    "            torch.nn.ReLU(),\n",
    "        )\n",
    "        self.downconv3 = torch.nn.Sequential(\n",
    "            torch.nn.Conv2d(256, 256, 3, 2, 1),\n",
    "            torch.nn.BatchNorm2d(256),\n",
    "            torch.nn.ReLU(),\n",
    "            \n",
    "            torch.nn.Conv2d(256, 512, 3, 1, 1),\n",
    "            torch.nn.BatchNorm2d(512),\n",
    "            torch.nn.ReLU(),\n",
    "            \n",
    "            torch.nn.Conv2d(512, 1024, 3, 1, 1),\n",
    "            torch.nn.BatchNorm2d(1024),\n",
    "            torch.nn.ReLU(),\n",
    "        )\n",
    "        self.flat = torch.nn.Sequential(\n",
    "            torch.nn.Conv2d(1024, 1024, 3, 1, 1),\n",
    "            torch.nn.BatchNorm2d(1024),\n",
    "            torch.nn.ReLU(),\n",
    "            \n",
    "            torch.nn.Conv2d(1024, 1024, 3, 1, 1),\n",
    "            torch.nn.BatchNorm2d(1024),\n",
    "            torch.nn.ReLU(),\n",
    "            \n",
    "            torch.nn.Conv2d(1024, 1024, 3, 1, 1),\n",
    "            torch.nn.BatchNorm2d(1024),\n",
    "            torch.nn.ReLU(),\n",
    "            \n",
    "            torch.nn.Conv2d(1024, 512, 3, 1, 1),\n",
    "            torch.nn.BatchNorm2d(512),\n",
    "            torch.nn.ReLU(),\n",
    "\n",
    "            torch.nn.Conv2d(512, 256, 3, 1, 1),\n",
    "            torch.nn.BatchNorm2d(256),\n",
    "            torch.nn.ReLU(),\n",
    "        )\n",
    "        \n",
    "        self.upconv1 = torch.nn.Sequential(\n",
    "#             torch.nn.Conv2d(256, 256, 4, 0.5, 1),\n",
    "            torch.nn.ConvTranspose2d(256, 256, 4, 2, 1),\n",
    "            torch.nn.BatchNorm2d(256),\n",
    "            torch.nn.ReLU(),\n",
    "            \n",
    "            torch.nn.Conv2d(256, 256, 3, 1, 1),\n",
    "            torch.nn.BatchNorm2d(256),\n",
    "            torch.nn.ReLU(),\n",
    "            \n",
    "            torch.nn.Conv2d(256, 128, 3, 1, 1),\n",
    "            torch.nn.BatchNorm2d(128),\n",
    "            torch.nn.ReLU(),\n",
    "        )\n",
    "        self.upconv2 = torch.nn.Sequential(\n",
    "#             torch.nn.Conv2d(128, 128, 4, 0.5, 1),\n",
    "            torch.nn.ConvTranspose2d(128, 128, 4, 2, 1),            \n",
    "            torch.nn.BatchNorm2d(128),\n",
    "            torch.nn.ReLU(),\n",
    "            \n",
    "            torch.nn.Conv2d(128, 128, 3, 1, 1),\n",
    "            torch.nn.BatchNorm2d(128),\n",
    "            torch.nn.ReLU(),\n",
    "            \n",
    "            torch.nn.Conv2d(128, 48, 3, 1, 1),\n",
    "            torch.nn.BatchNorm2d(48),\n",
    "            torch.nn.ReLU(),\n",
    "        )\n",
    "        self.upconv3 = torch.nn.Sequential(\n",
    "#             torch.nn.Conv2d(48, 48, 4, 0.5, 1),\n",
    "            torch.nn.ConvTranspose2d(48, 48,4, 2, 1),\n",
    "            torch.nn.BatchNorm2d(48),\n",
    "            torch.nn.ReLU(),\n",
    "            \n",
    "            torch.nn.Conv2d(48, 24, 3, 1, 1),\n",
    "            torch.nn.BatchNorm2d(24),\n",
    "            torch.nn.ReLU(),\n",
    "            \n",
    "            torch.nn.Conv2d(24, 1, 3, 1, 1),\n",
    "            torch.nn.Sigmoid (),\n",
    "        )\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        conv1_out = self.downconv1(x)\n",
    "        conv2_out = self.downconv2(conv1_out)\n",
    "        conv3_out = self.downconv3(conv2_out)\n",
    "        flat_out = self.flat(conv3_out)\n",
    "        upconv1_out = self.upconv1(flat_out)\n",
    "        upconv2_out = self.upconv2(upconv1_out)\n",
    "        upconv3_out = self.upconv3(upconv2_out)\n",
    "        return upconv3_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LossFunc(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super(LossFunc,self).__init__()\n",
    "        return\n",
    "    \n",
    "    def forward(self,pred,gt):\n",
    "        dis = torch.abs(pred-gt)\n",
    "        sqr = torch.pow(dis,2).sum()\n",
    "#         loss = torch.pow(sqr,0.5)\n",
    "        size = pred.size()\n",
    "        all_pix_c = 1\n",
    "        for s in size:\n",
    "            all_pix_c = all_pix_c*s \n",
    "        loss = sqr / all_pix_c\n",
    "        return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from time import time\n",
    "model = torch.nn.DataParallel(Net())\n",
    "optimizer = torch.optim.Adadelta(model.parameters())\n",
    "loss_func = LossFunc()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in range(3):\n",
    "    # training-----------------------------\n",
    "    train_loss = 0.\n",
    "    train_acc = 0.\n",
    "    start = time()\n",
    "    for batch_x, batch_y in train_loader:\n",
    "#         print(batch_x)\n",
    "#         print(batch_y)\n",
    "        batch_x, batch_y = Variable(batch_x).cuda(), Variable(batch_y).cuda()\n",
    "        out = model(batch_x)\n",
    "#         print(batch_x.size())\n",
    "#         print(out.size())\n",
    "        loss = loss_func(out, batch_y)\n",
    "        train_loss += loss.data[0]\n",
    "        pred = out\n",
    "#         print(pred)\n",
    "#         print(batch_y)\n",
    "        size = pred.size()\n",
    "        all_pix = 1\n",
    "        for s in size:\n",
    "            all_pix = all_pix*s \n",
    "        acc = (torch.abs(pred - batch_y)>0.01).sum()/all_pix\n",
    "        train_acc +=acc.data[0]\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "    print(\"Epoch %d. Loss: %.6f, Train acc %.2f, Time %.1f sec\" % (\n",
    "            epoch+1, train_loss / (len(train_data)), train_acc / len(train_data), time() - start\n",
    "        ))\n",
    "    \n",
    "    if epoch % 1 ==0:\n",
    "#         PATH = 'model/epoch-'+str(epoch)\n",
    "#         torch.save(model, PATH)\n",
    "\n",
    "        idx = 0\n",
    "        for batch_x, batch_y in test_loader:\n",
    "            idx = idx+1\n",
    "            batch_x, batch_y = Variable(batch_x).cuda(), Variable(batch_y).cuda()\n",
    "            try:\n",
    "                start = time()\n",
    "                out = model(batch_x)\n",
    "                pred = out\n",
    "                for i in range(len(batch_x)):\n",
    "                    predimg = (pred[i].cpu().data.numpy() * 255)[0,:,:]\n",
    "                    pilimg = Image.fromarray(np.uint8(predimg))\n",
    "                    pilimg.save('pred/'+str(epoch)  + str(idx)+'_pred.png')             \n",
    "            except:\n",
    "                continue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "test_data=MyDataset(root+'train2.txt', -1, 1,transform=transforms.ToTensor())\n",
    "test_loader = DataLoader(dataset=test_data, batch_size=1)\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "idx = 0\n",
    "for batch_x, batch_y in test_loader:\n",
    "    batch_x, batch_y = Variable(batch_x).cuda(), Variable(batch_y).cuda()\n",
    "    try:\n",
    "        start = time()\n",
    "        out = model(batch_x)\n",
    "        pred = out\n",
    "        print('%.6f sec'%(time()-start))\n",
    "        _,figs = plt.subplots(1,3)\n",
    "        for i in range(len(batch_x)):\n",
    "            img = (batch_x[i].cpu().data.numpy() * 255)[0,:,:]\n",
    "            predimg = (pred[i].cpu().data.numpy() * 255)[0,:,:]\n",
    "            gtimg = (batch_y[i].cpu().data.numpy() * 255)[0,:,:]\n",
    "            figs[0].imshow(img,cmap = 'gray')\n",
    "            figs[1].imshow(predimg,cmap = 'gray')\n",
    "            figs[2].imshow(gtimg,cmap = 'gray')\n",
    "\n",
    "            pilimg = Image.fromarray(np.uint8(img))\n",
    "            pilimg.save('pred/' + str(idx)+'.png')\n",
    "\n",
    "            pilimg = Image.fromarray(np.uint8(predimg))\n",
    "            pilimg.save('pred/' + str(idx)+'_pred.png')\n",
    "\n",
    "            pilimg = Image.fromarray(np.uint8(gtimg))\n",
    "            pilimg.save('pred/' + str(idx)+'_gt.png')\n",
    "        plt.show()\n",
    "        idx = idx+1\n",
    "    except:\n",
    "        print('Out of Memory!')\n",
    "        continue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}