{"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":"# from google.colab import drive\n# drive.mount('/content/drive')","metadata":{"id":"8mGy53qr0dDo","outputId":"e2516ceb-0757-4b01-b3f1-9b1b7f15bb90","execution":{"iopub.status.busy":"2022-03-20T16:51:24.279177Z","iopub.execute_input":"2022-03-20T16:51:24.279538Z","iopub.status.idle":"2022-03-20T16:51:24.286633Z","shell.execute_reply.started":"2022-03-20T16:51:24.279488Z","shell.execute_reply":"2022-03-20T16:51:24.285871Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Consts","metadata":{"id":"03_I7zlDtgOH"}},{"cell_type":"code","source":"COLAB=False\nDIR='./' #'drive/MyDrive/deep_learning_project'\nMASK_DIR='../input/deep-learning-project/masks/test_center_masks'\nMODEL='auto_encoder' #'UNET_WITH_CONTENT_LOSS' #'UNET' #\nMASK='center'\nCONTENT_LOSS=False\nSTYLE_LOSS=False","metadata":{"id":"NR8fKQxxthdK","execution":{"iopub.status.busy":"2022-03-20T16:51:24.288451Z","iopub.execute_input":"2022-03-20T16:51:24.289404Z","iopub.status.idle":"2022-03-20T16:51:24.297545Z","shell.execute_reply.started":"2022-03-20T16:51:24.289348Z","shell.execute_reply":"2022-03-20T16:51:24.296532Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install pytorch-ignite","metadata":{"id":"fDAr6mHZm46N","outputId":"20fcef0b-d20b-4402-ad29-69fb5730f258","execution":{"iopub.status.busy":"2022-03-20T16:51:24.299078Z","iopub.execute_input":"2022-03-20T16:51:24.300064Z","iopub.status.idle":"2022-03-20T16:51:31.837388Z","shell.execute_reply.started":"2022-03-20T16:51:24.300022Z","shell.execute_reply":"2022-03-20T16:51:31.836541Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# ! unzip -n drive/MyDrive/deep_learning_project/data/data.zip\n# ! unzip -n drive/MyDrive/deep_learning_project/data/masks.zip\n# ! unzip -n drive/MyDrive/deep_learning_project/data/monet_data_set.zip\n# ! unzip -n drive/MyDrive/deep_learning_project/data/test.zip","metadata":{"id":"hJubwcx30Zye","outputId":"7f9870a8-4ac8-4d68-c69b-47e70ccdee3c","execution":{"iopub.status.busy":"2022-03-20T16:51:31.841501Z","iopub.execute_input":"2022-03-20T16:51:31.841728Z","iopub.status.idle":"2022-03-20T16:51:31.846106Z","shell.execute_reply.started":"2022-03-20T16:51:31.841700Z","shell.execute_reply":"2022-03-20T16:51:31.845402Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"imports","metadata":{"id":"A32PZwO80J-A"}},{"cell_type":"code","source":"import math\nimport numbers\nimport os\nimport random\nimport sys\nimport warnings\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport torch\nfrom torch import Tensor\nimport torchvision\nfrom IPython.display import Image\nfrom PIL import Image\nfrom scipy import ndimage\nfrom ignite.engine import Engine\nfrom ignite.metrics import PSNR, MeanSquaredError, Loss\nfrom torch import nn\nfrom torch.utils.data import DataLoader\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms,models\n# from torchvision.transforms import functional as F\nfrom torch.nn import functional as F\nfrom tqdm.notebook import tqdm\nfrom torch.cuda import amp\nimport glob\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(f\"Using {device} device\")","metadata":{"id":"u3ymWald0J-C","outputId":"3ff4395c-823f-4f8b-c216-85be7ae6a2b6","execution":{"iopub.status.busy":"2022-03-20T16:51:31.847281Z","iopub.execute_input":"2022-03-20T16:51:31.847828Z","iopub.status.idle":"2022-03-20T16:51:31.861035Z","shell.execute_reply.started":"2022-03-20T16:51:31.847786Z","shell.execute_reply":"2022-03-20T16:51:31.860171Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Dataset loader","metadata":{"id":"oWLWOLg-0J-D"}},{"cell_type":"code","source":"class ImageTestDataSet(Dataset):\n    def __init__(self,data_dir,masks_dir,transforms=None,monet=False):\n        images_dir = os.path.join(data_dir, masks_dir)\n        masks = glob.glob(f\"{images_dir}/*mask*.jpg\")\n        images = glob.glob(f\"{images_dir}/*result*.jpg\")\n        self.files = list(zip(images, masks))\n        self.transforms = transforms\n\n    def __len__(self):\n        return len(self.files)\n\n    def __getitem__(self, index):\n        img,mask = self.files[index]\n        img = Image.open(img)\n        mask = mask_transformer()(Image.open(mask))\n        if self.transforms is not None:\n            img = self.transforms(img)\n\n        return img, mask\n    \nclass ImageDataset(Dataset):\n    def __init__(self, data_dir, mode='validate', transforms=None, monet=False,mask = 'center'):\n        images_dir = os.path.join(data_dir, 'monet_jpg' if monet else 'photo_jpg')\n        image_list = os.listdir(images_dir)\n        train_size = int(0.80 * len(image_list))\n        if mode == 'train':\n            self.files = [os.path.join(images_dir, name) for name in sorted(os.listdir(images_dir))[:train_size]]\n        elif mode == 'validate':\n            self.files = [os.path.join(images_dir, name) for name in sorted(os.listdir(images_dir))[train_size:]]\n        if MASK_DIR:\n            self.mask_function = MaskFromImages(MASK_DIR)\n        else:\n            if mask == 'center':\n                self.mask_function = central_region_transformer()\n            elif mask == 'block':\n                self.mask_function = random_blocks_transformer()\n            else:\n                self.mask_function = MaskFromImages()\n        self.transforms = transforms\n\n    def __len__(self):\n        return len(self.files)\n\n    def __getitem__(self, index):\n        file = self.files[index]\n        img = Image.open(file)\n\n        if self.transforms is not None:\n            original_img = self.transforms[0](img)\n            mask = self.mask_function(torch.zeros_like(original_img))\n            img_with_erased_region = torch.where(mask == 1, mask, original_img)\n            erased_region = 1 - abs(img_with_erased_region - original_img)\n            return original_img, erased_region, img_with_erased_region, mask\n\n        return img, img, img,img\n\n\n","metadata":{"id":"1mOYyihO0J-D","execution":{"iopub.status.busy":"2022-03-20T16:51:31.862672Z","iopub.execute_input":"2022-03-20T16:51:31.863300Z","iopub.status.idle":"2022-03-20T16:51:31.881778Z","shell.execute_reply.started":"2022-03-20T16:51:31.863257Z","shell.execute_reply":"2022-03-20T16:51:31.880983Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Models","metadata":{"id":"zHRmUCAJ0J-E"}},{"cell_type":"code","source":"class AutoEncoder(nn.Module):\n    def __init__(self, encoder=None, decoder=None):\n        super().__init__()\n        if encoder is None or decoder is None:\n            self.decoder, self.encoder = auto_encoder_parameters()\n        else:\n            self.encoder = encoder\n            self.decoder = decoder\n\n    def forward(self, x):\n        encoder_output = self.encoder(x)\n\n        output = self.decoder(encoder_output)\n        return output\n\n\nclass Discriminator(nn.Module):\n    def __init__(self, model=None):\n        super().__init__()\n        self.model = discriminator_parameters() if model is None else model\n\n    def forward(self, x):\n        output = self.model(x)\n        return output\n","metadata":{"id":"1L1HBJZ50J-E","execution":{"iopub.status.busy":"2022-03-20T16:51:31.883454Z","iopub.execute_input":"2022-03-20T16:51:31.883965Z","iopub.status.idle":"2022-03-20T16:51:31.894743Z","shell.execute_reply.started":"2022-03-20T16:51:31.883880Z","shell.execute_reply":"2022-03-20T16:51:31.893933Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Util","metadata":{"id":"l1DZBJwD0J-E"}},{"cell_type":"code","source":"from matplotlib.rcsetup import validate_fontsize\ndef crop_image(batch):\n    batch_size, dim, img_w, img_h = batch.shape\n    i = (img_h - 64) // 2\n    j = (img_w - 64) // 2\n    croped_batch = batch[:, :, i:i + 64, j:j + 64]\n    return croped_batch\n\n\ndef save_weights(model, path):\n    torch.save(model.state_dict(), path)\n\n\ndef load_weights(model, path):\n    model.load_state_dict(torch.load(path))\n\n\ndef add_missing_region_to_image(img, region):\n    dim, img_w, img_h = img.shape\n    i = (img_h - 64) // 2\n    j = (img_w - 64) // 2\n    img[:, i:i + 64, j:j + 64] = region\n    return img\n\n\ndef get_image_for_plot(fake, img_with_erased_region, mask, original_img, real):\n    index = random.randint(0, real.shape[0] - 1)\n    full_pred = fake[index].detach().cpu()\n    mask = mask.detach().cpu()\n    image_after_mask = img_with_erased_region[index]\n    full_pred = torch.where(mask[index] == 1, full_pred, image_after_mask)\n    pred = torch.where(mask[index] == 0, torch.ones_like(full_pred), full_pred)\n    pred = torch.where(mask[index] == 1, full_pred, pred)\n    full_pred = full_pred.permute(1, 2, 0).detach().cpu()\n    pred = pred.permute(1, 2, 0).detach().cpu()\n    y = original_img[index].permute(1, 2, 0).detach().cpu()\n    image_after_mask = image_after_mask.permute(1, 2, 0).detach().cpu()\n    return full_pred, pred, y, image_after_mask\n\n\ndef show_images(train_true, train_fake, validate_true, validate_fake):\n    train_true = (1 + train_true) / 2\n    train_fake = (1 + train_fake) / 2\n    validate_true = (1 + validate_true) / 2\n    validate_fake = (1 + validate_fake) / 2\n    fig, axs = plt.subplots(2, 2)\n    fig.set_dpi(120)\n    axs[0, 0].imshow(train_true, aspect='auto')\n    axs[0, 0].set(ylabel='true')\n    axs[0, 0].axis('off')\n    axs[1, 0].imshow(train_fake, aspect='auto')\n    axs[1, 0].set(xlabel='train', ylabel='fake')\n    axs[1, 0].axis('off')\n    axs[0, 1].imshow(validate_true, aspect='auto')\n    axs[0, 1].axis('off')\n    axs[1, 1].imshow(validate_fake, aspect='auto')\n    axs[1, 1].set(xlabel='validate')\n    axs[1, 1].axis('off')\n    axs[0, 0].set_title('Train')\n    axs[0, 1].set_title('Validate')\n    plt.subplots_adjust(hspace=0, wspace=0)\n    plt.tight_layout()\n    plt.show()\n\ndef show_with_missing_regions(train_full_fake, train_true, train_fake, train_masked,\n                              validate_full_fake, validate_true,\n                              validate_fake, validate_masked):\n    train_fake, train_full_fake, train_masked, train_true, validate_fake, validate_full_fake, validate_masked, validate_true =\\\n    denormalize_images(train_full_fake, train_masked, train_fake,\n    train_true, validate_full_fake,validate_masked, validate_fake, validate_true)\n    fig, axs = plt.subplots(2, 4)\n    fig.set_dpi(120)\n    axs[0, 0].imshow(np.clip(train_true,0,1))\n    axs[0, 1].imshow(np.clip(train_masked,0,1))\n    axs[0, 2].imshow(np.clip(train_fake,0,1))\n    axs[0, 3].imshow(np.clip(train_full_fake,0,1))\n    axs[1, 0].imshow(np.clip(validate_true,0,1))\n    axs[1, 1].imshow(np.clip(validate_masked,0,1))\n    axs[1, 2].imshow(np.clip(validate_fake,0,1))\n    axs[1, 3].imshow(np.clip(validate_full_fake,0,1))\n\n    for a in axs:\n        for x in a:\n            x.axis('off')\n\n    axs[0, 0].set_title('Train', fontsize=16)\n    axs[1, 0].set_title('Validate', fontsize=16)\n    plt.subplots_adjust(hspace=0, wspace=0)\n    plt.tight_layout()\n    plt.show()\n\n\ndef denormalize_images(train_full_fake, train_image_after_mask, train_fake, train_true, validate_full_fake,\n                       validate_masked, validate_fake, validate_true):\n    train_true = (1 + train_true) / 2\n    train_masked = (1 + train_image_after_mask) / 2\n    train_fake = (1 + train_fake) / 2\n    train_full_fake = (1 + train_full_fake) / 2\n    validate_true = (1 + validate_true) / 2\n    validate_masked = (1 + validate_masked) / 2\n    validate_fake = (1 + validate_fake) / 2\n    validate_full_fake = (1 + validate_full_fake) / 2\n    return train_fake, train_full_fake, train_masked, train_true, validate_fake, validate_full_fake, validate_masked, validate_true\n\n\n\ndef show_sample(data_set):\n    img_a, img_b, img_c = data_set[0]\n    plt.imshow(img_a.permute(1, 2, 0))\n    plt.show()\n    plt.imshow(img_b.permute(1, 2, 0))\n    plt.show()\n    plt.imshow(img_c.permute(1, 2, 0))\n    plt.show()\n\n\n\ndef get_data_from_files(data_dir='./'):\n    train_dataset = ImageDataset(data_dir=data_dir, mode='train',\n                                 transforms=[data_augmentation_transformer()],mask=MASK)\n    # show_sample(train_dataset)\n    validate_dataset = ImageDataset(data_dir=data_dir, mode='validate',\n                                    transforms=[general_transformer()],mask=MASK)\n    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)\n    validate_loader = DataLoader(validate_dataset, batch_size=64, shuffle=True)\n\n    return train_loader, validate_loader\n\ndef weights_init(m):\n    classname = m.__class__.__name__\n    if classname.find('Conv') != -1:\n        nn.init.normal_(m.weight.data, 0.0, 0.02)\n    elif classname.find('BatchNorm') != -1:\n        nn.init.normal_(m.weight.data, 1.0, 0.02)\n        nn.init.constant_(m.bias.data, 0)\n\ndef save_models(ae_pack,discriminator_pack,loss,epoch,evaluator):\n  print(f'--- save {evaluator} models ---')\n  torch.save({\n            'epoch': epoch,\n            'ae_state_dict': ae_pack['model'].state_dict(),\n            'ae_optimizer_state_dict': ae_pack['optimizer'].state_dict(),\n            'd_state_dict': discriminator_pack['model'].state_dict(),\n            'd_optimizer_state_dict': discriminator_pack['optimizer'].state_dict(),\n            'loss': loss,\n            \n            }, f'{DIR}/{MODEL}_{MASK}_{evaluator}')\n","metadata":{"id":"VK-EM4V30J-F","execution":{"iopub.status.busy":"2022-03-20T16:51:31.897683Z","iopub.execute_input":"2022-03-20T16:51:31.897923Z","iopub.status.idle":"2022-03-20T16:51:31.935801Z","shell.execute_reply.started":"2022-03-20T16:51:31.897889Z","shell.execute_reply":"2022-03-20T16:51:31.934930Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Masks","metadata":{"id":"AEyF19DS0J-G"}},{"cell_type":"code","source":"from torchvision.transforms import functional as F1\nclass MaskFromImages:\n    def __init__(self, dir='deep-learning-project/masks/masks'):\n        self.files = [os.path.join(dir, name) for name in sorted(os.listdir(dir))]\n\n    def __call__(self,img):\n        index = np.random.randint(len(self.files))\n        mask = self.files[index]\n        mask = Image.open(mask)\n        mask = transforms.ToTensor()(mask)\n        img = torch.where(mask == 1,mask,img )\n        return img\n\n\n\n\n\nclass RandomBlocksErasing:\n\n    def __init__(self, p=0.5, scale=(0.02, 0.33), ratio=(0.3, 3.3), value=0, center=True, inplace=False,\n                 random_blocks=0):\n        assert isinstance(value, (numbers.Number, str, tuple, list))\n        if (scale[0] > scale[1]) or (ratio[0] > ratio[1]):\n            warnings.warn(\"range should be of kind (min, max)\")\n        if scale[0] < 0 or scale[1] > 1:\n            raise ValueError(\"range of scale should be between 0 and 1\")\n        if p < 0 or p > 1:\n            raise ValueError(\"range of random erasing probability should be between 0 and 1\")\n\n        self.p = p\n        self.scale = scale\n        self.ratio = ratio\n        self.value = value\n        self.inplace = inplace\n        self.center = center\n        self.random_blocks = random_blocks\n\n    @staticmethod\n    def get_params(img, scale, ratio, value=0, center=True):\n        \"\"\"Get parameters for ``erase`` for a random erasing.\n        Args:\n            img (Tensor): Tensor image of size (C, H, W) to be erased.\n            scale: range of proportion of erased area against input image.\n            ratio: range of aspect ratio of erased area.\n        Returns:\n            tuple: params (i, j, h, w, v) to be passed to ``erase`` for random erasing.\n        \"\"\"\n        img_c, img_h, img_w = img.shape\n        area = img_h * img_w\n        while True:\n            erase_area = random.uniform(scale[0], scale[1]) * area\n            aspect_ratio = random.uniform(ratio[0], ratio[1])\n\n            h = int(round(math.sqrt(erase_area * aspect_ratio)))\n            w = int(round(math.sqrt(erase_area / aspect_ratio)))\n\n            if h < img_h and w < img_w:\n                if center:\n                    i = (img_h - h) // 2\n                    j = (img_w - w) // 2\n                else:\n                    i = random.randint(0, img_h - h)\n                    j = random.randint(0, img_w - w)\n                if isinstance(value, numbers.Number):\n                    v = value\n                elif isinstance(value, torch._six.string_classes):\n                    v = torch.empty([img_c, h, w], dtype=torch.float32).normal_()\n                elif isinstance(value, (list, tuple)):\n                    v = torch.tensor(value, dtype=torch.float32).view(-1, 1, 1).expand(-1, h, w)\n                else:\n                    v = None\n                return i, j, h, w, v\n\n        # Return original image\n        return 0, 0, img_h, img_w, img\n\n    def __call__(self, img):\n        \"\"\"\n        Args:\n            img (Tensor): Tensor image of size (C, H, W) to be erased.\n        Returns:\n            img (Tensor): Erased Tensor image.\n        \"\"\"\n        if self.random_blocks > 0:\n            num_of_blocks = random.randint(1, self.random_blocks)\n            for i in range(num_of_blocks):\n                if random.uniform(0, 1) < self.p:\n                    x, y, h, w, v = self.get_params(img, scale=self.scale, ratio=self.ratio, value=self.value,\n                                                    center=False)\n                    img = F1.erase(img, x, y, h, w, v, self.inplace)\n            return img\n        else:\n            if random.uniform(0, 1) < self.p:\n                x, y, h, w, v = self.get_params(img, scale=self.scale, ratio=self.ratio, value=self.value,\n                                                center=self.center)\n                return F1.erase(img, x, y, h, w, v, self.inplace)\n            return img\n\n\nclass RandomRegionErasing:\n    def __init__(self, p=0.5, region_size=100, inplace=False):\n        self.p = p\n        self.region_size = region_size\n        self.inplace = inplace\n\n    @staticmethod\n    def get_random_mask(img, region_size):\n        img_c, img_h, img_w = img.shape\n\n        n = 10\n        mask = np.zeros((region_size, region_size))\n        generator = np.random.RandomState()\n        points = region_size * generator.rand(2, n ** 2)\n        mask[(points[0]).astype(np.int), (points[1]).astype(np.int)] = 1\n        mask = ndimage.gaussian_filter(mask, sigma=region_size / (4. * n))\n        mask = (mask > mask.mean()).astype(np.float)\n        img = np.ones((img_h, img_w))\n        start_h, start_w = random.randint(0, img_w - region_size), random.randint(0, img_w - region_size)\n\n        img[start_h:start_h + region_size, start_w:start_w + region_size] = mask\n        return img\n\n\n    def __call__(self, img):\n        \"\"\"\n        Args:\n            img (Tensor): Tensor image of size (C, H, W) to be erased.\n        Returns:\n            img (Tensor): Erased Tensor image.\n        \"\"\"\n        if random.uniform(0, 1) < self.p:\n            mask = self.get_random_mask(img, region_size=self.region_size)\n\n            if not self.inplace:\n                img = img.clone()\n            mask_3d = mask[None, :, :] * np.ones(3, dtype=int)[:, None, None]\n            indices_mask = np.where(mask_3d == 0)\n            img[indices_mask] = 1\n\n        return img\n","metadata":{"id":"5ZddcCX20J-G","execution":{"iopub.status.busy":"2022-03-20T16:51:31.939217Z","iopub.execute_input":"2022-03-20T16:51:31.939433Z","iopub.status.idle":"2022-03-20T16:51:31.970358Z","shell.execute_reply.started":"2022-03-20T16:51:31.939405Z","shell.execute_reply":"2022-03-20T16:51:31.969542Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Transforms","metadata":{"id":"c0HDVBXD0J-H"}},{"cell_type":"code","source":"def central_region_transformer(p=1, scale=(0.0625, 0.0625), value=1, ratio=(1, 1)):\n    return transforms.Compose([\n        RandomBlocksErasing(p, scale, ratio, value, center=True, random_blocks=0),\n    ])\n\n\ndef random_blocks_transformer(p=1, scale=(0.02, 0.02), value=1, ratio=(1, 1)):\n    return transforms.Compose([\n        RandomBlocksErasing(p, scale, ratio, value, center=False, random_blocks=10),\n        \n\n    ])\n\n\ndef random_region_transformer(p=1, region_size=100):\n    return transforms.Compose([\n        RandomRegionErasing(p, region_size=region_size),\n    ])\n\n\ndef data_augmentation_transformer():\n    return transforms.Compose([\n        transforms.RandomHorizontalFlip(p=0.5),\n        transforms.RandomRotation(degrees=(-8, 8)),\n#         transforms.Resize((224, 224)),\n        #         transforms.RandomResizedCrop(size=(256, 256)),\n        # transforms.Resize((128,128)),\n        #         transforms.RandomAdjustSharpness(sharpness_factor=1, p=0.2),\n        transforms.ToTensor(),\n        # RandomBlocksErasing(p=0.5, scale=(0.02, 0.02), value=1, ratio=(1, 1), random_blocks=10),\n        transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))\n    ])\ndef mask_transformer():\n    return transforms.Compose([\n         transforms.Resize((256, 256)),\n        #         transforms.RandomResizedCrop(size=(256, 256)),\n        # transforms.Resize((128,128)),\n        #         transforms.RandomAdjustSharpness(sharpness_factor=1, p=0.2),\n        transforms.ToTensor(),\n        # RandomBlocksErasing(p=0.5, scale=(0.02, 0.02), value=1, ratio=(1, 1), random_blocks=10),\n#         transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))\n    ])\n\n\ndef general_transformer():\n    return torchvision.transforms.Compose([\n#         transforms.Resize((224, 224)),\n        transforms.ToTensor(),\n        transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))\n    ])\n","metadata":{"id":"IZLkRlnC0J-H","execution":{"iopub.status.busy":"2022-03-20T16:51:31.973601Z","iopub.execute_input":"2022-03-20T16:51:31.974146Z","iopub.status.idle":"2022-03-20T16:51:31.988080Z","shell.execute_reply.started":"2022-03-20T16:51:31.974103Z","shell.execute_reply":"2022-03-20T16:51:31.987228Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"train and validate","metadata":{"id":"ImZil2Ug0J-I"}},{"cell_type":"code","source":"\n\ndef train(dataloader, auto_encoder_pack, discriminator_pack):\n    auto_encoder, ae_criterion, ae_optimizer, a = auto_encoder_pack['model'], auto_encoder_pack['loss'], \\\n                                                  auto_encoder_pack[\n                                                      'optimizer'], auto_encoder_pack['lambda']\n    discriminator, discriminator_criterion, discriminator_optimizer = discriminator_pack['model'], discriminator_pack[\n        'loss'], discriminator_pack['optimizer']\n\n    real_label, fake_label = 1, 0\n    size = len(dataloader.batch_sampler)\n    pbar = tqdm(range(size))\n    auto_encoder.train()\n    discriminator.train()\n    train_ae_loss, train_d_loss, correct = 0, 0, 0\n\n    for batch, X in zip(pbar, dataloader):\n\n        original_img, erased_region, img_with_erased_region, mask = X\n        X = img_with_erased_region.to(device)\n        mask = mask.to(device)\n\n        real = original_img.to(device)\n        b_size = real.size(0)\n        r_label = torch.full((b_size,), real_label, dtype=torch.float, device=device)\n        f_label = torch.full((b_size,), fake_label, dtype=torch.float, device=device)\n\n        # -------------------\n        # Train Auto Encoder\n        # -------------------\n\n        auto_encoder.zero_grad()\n        ae_optimizer.zero_grad()\n\n        fake = auto_encoder(X)\n        erased_region = erased_region.to(device)\n        generated_region = torch.where(mask == 1, fake, erased_region)\n        f_output = discriminator(fake.detach()).view(-1)\n        adversarial_loss = discriminator_criterion(f_output, r_label)  # + discriminator_criterion(r_output,r_label)\n        errAE = ae_criterion(fake, real) + ae_criterion(erased_region, generated_region)\n        loss = a * errAE + (1 - a) * (adversarial_loss)\n        train_ae_loss += loss.item()\n        loss.backward()\n        ae_optimizer.step()\n\n        # -------------------\n        # Train Discriminator\n        # -------------------\n\n        discriminator.zero_grad()\n        discriminator_optimizer.zero_grad()\n\n        # --- Real Batch ---\n\n        # Forward\n        r_output = discriminator(real).view(-1)\n        correct += (r_output > 0.5).type(torch.float).sum().item()\n\n        # Calculate loss\n        errD_real = discriminator_criterion(r_output, r_label)\n        # --- Fake Batch ---\n\n        # Forward\n        f_output = discriminator(fake.detach()).view(-1)\n        correct += (f_output < 0.5).type(torch.float).sum().item()\n\n        # Calculate loss\n        errD_fake = discriminator_criterion(f_output, f_label)\n\n        # Compute error of D as sum over the fake and the real batches\n        errD = 0.5 * (errD_real + errD_fake)\n\n        # Backward\n        errD.backward()\n\n        # Update D\n        discriminator_optimizer.step()\n\n        train_d_loss += errD.item()\n\n        pbar.set_description(\n            f\"Batch {batch} | AutoEncoder loss = {train_ae_loss / (batch + 1):>7f} Discriminator loss = {train_d_loss / (batch + 1):>7f}\")\n\n    train_ae_loss /= len(dataloader)\n    train_d_loss /= len(dataloader)\n    correct *= (50 / (len(dataloader.dataset)))\n    full_pred, pred, y, image_after_mask = get_image_for_plot(fake, img_with_erased_region, mask, original_img, real)\n\n    print(f\"acc:{correct:.3f}\")\n    return train_ae_loss, full_pred, y, pred, image_after_mask\n\n\ndef train_with_cuda(dataloader, auto_encoder_pack, discriminator_pack, scaler):\n    auto_encoder, ae_criterion, ae_optimizer, a, content_criterion = auto_encoder_pack['model'], auto_encoder_pack[\n        'loss'], \\\n                                                                     auto_encoder_pack['optimizer'], auto_encoder_pack[\n                                                                         'lambda'], auto_encoder_pack['content_loss']\n    discriminator, discriminator_criterion, discriminator_optimizer = discriminator_pack['model'], discriminator_pack[\n        'loss'], discriminator_pack['optimizer']\n\n    real_label, fake_label = 1, 0\n    size = len(dataloader.batch_sampler)\n    pbar = tqdm(range(size))\n    auto_encoder.train()\n    discriminator.train()\n    train_ae_loss, train_d_loss, correct = 0, 0, 0\n\n    for batch, X in zip(pbar, dataloader):\n        # torch.cuda.empty_cache()\n\n        original_img, erased_region, img_with_erased_region, mask = X\n        X = img_with_erased_region.to(device)\n        mask = mask.to(device)\n\n        real = original_img.to(device)\n        b_size = real.size(0)\n        r_label = torch.full((b_size,), real_label, dtype=torch.float, device=device)\n        f_label = torch.full((b_size,), fake_label, dtype=torch.float, device=device)\n\n        # -------------------\n        # Train Auto Encoder\n        # -------------------\n\n        # auto_encoder.zero_grad()\n        ae_optimizer.zero_grad()\n        fake = auto_encoder(X)\n        erased_region = erased_region.to(device)\n        generated_region = torch.where(mask == 1, fake, erased_region)\n        with amp.autocast():\n            with torch.no_grad():\n                f_output = discriminator(fake.detach()).view(-1)\n            pixel_loss = ae_criterion(fake, real) + ae_criterion(erased_region, generated_region)\n            content_loss = content_criterion(fake, real) if CONTENT_LOSS else 0\n            adversarial_loss = discriminator_criterion(f_output, r_label)\n\n        # f_output = discriminator(fake.detach()).view(-1)\n        # adversarial_loss = discriminator_criterion(f_output, r_label)\n        # errAE = ae_criterion(fake, real) + ae_criterion(erased_region,generated_region)\n        loss = a * pixel_loss + (1 - a) * (adversarial_loss) + content_loss\n        train_ae_loss += loss.item()\n\n        scaler.scale(loss).backward()\n        # Update generator parameters\n        scaler.step(ae_optimizer)\n        scaler.update()\n\n        # loss.backward()\n        # ae_optimizer.step()\n\n        # -------------------\n        # Train Discriminator\n        # -------------------\n\n        # discriminator.zero_grad()\n        discriminator_optimizer.zero_grad()\n\n        # --- Real Batch ---\n\n        # Forward\n        r_output = discriminator(real).view(-1)\n        correct += (r_output > 0.5).type(torch.float).sum().item()\n\n        # Calculate loss\n        errD_real = discriminator_criterion(r_output, r_label)\n        # --- Fake Batch ---\n\n        # Forward\n        f_output = discriminator(fake.detach()).view(-1)\n        correct += (f_output < 0.5).type(torch.float).sum().item()\n\n        # Calculate loss\n        errD_fake = discriminator_criterion(f_output, f_label)\n\n        # Compute error of D as sum over the fake and the real batches\n        errD = 0.5 * (errD_real + errD_fake)\n\n        # Backward\n        # errD.backward()\n        scaler.scale(errD).backward()\n        scaler.step(discriminator_optimizer)\n        scaler.update()\n\n        # Update D\n        # discriminator_optimizer.step()\n\n        train_d_loss += errD.item()\n\n        pbar.set_description(\n            f\"Batch {batch} | AutoEncoder loss = {train_ae_loss / (batch + 1):>7f} Discriminator loss = {train_d_loss / (batch + 1):>7f}\")\n\n    train_ae_loss /= len(dataloader)\n    train_d_loss /= len(dataloader)\n    correct *= (50 / (len(dataloader.dataset)))\n    full_pred, pred, y, image_after_mask = get_image_for_plot(fake, img_with_erased_region, mask, original_img, real)\n\n    print(f\"acc:{correct:.3f}\")\n    return train_ae_loss, full_pred, y, pred, image_after_mask\n\n\n\n\ndef validate(dataloader, model, loss_fn, check_best=False):\n    global best\n\n    def predict_on_batch(engine, batch):\n        with torch.no_grad():\n            original_img, erased_region, img_with_erased_region, mask = batch\n            X = img_with_erased_region.to(device)\n            y = original_img.to(device)\n            fake = model(X)\n\n        return fake, y\n\n    evaluator = Engine(predict_on_batch)\n    psnr = PSNR(1.0,device=device)\n    psnr.attach(evaluator, 'psnr')\n    mse = Loss(nn.MSELoss())\n    mse.attach(evaluator,'mse')\n    state = evaluator.run(dataloader)\n    validate_mse = state.metrics['mse']\n    validate_psnr= state.metrics['psnr']\n    with torch.no_grad():\n        for X in dataloader:\n            original_img, erased_region, img_with_erased_region, mask = X\n            X = img_with_erased_region.to(device)\n            fake = model(X)\n            break\n\n    full_pred, pred, y, image_after_mask = get_image_for_plot(fake, img_with_erased_region, mask, original_img,\n                                                              original_img)\n\n    print(f\"Validate Error:\\n MSE loss: {validate_mse:>5f} PSNR: {validate_psnr:>5f}\\n\")\n    return validate_mse, validate_psnr, full_pred, y, pred, image_after_mask\n\n\n#\n# def validate(dataloader, model, loss_fn, check_best=False):\n#     global best\n#     size = len(dataloader.dataset)\n#     num_batches = len(dataloader)\n#     psnr = PSNR(1.0,device=device)\n#     # psnr.attach(loss_fn, 'psnr')\n#     model.eval()\n#     validate_loss, correct = 0, 0\n#     with torch.no_grad():\n#         for X in dataloader:\n#             original_img, erased_region, img_with_erased_region, mask = X\n#             X = img_with_erased_region.to(device)\n#             y = original_img.to(device)\n#             fake = model(X)\n#             validate_loss = psnr.run(fake, y)\n#             break\n#\n#     print(validate_loss.metrics['psnr'])\n#     validate_loss /= num_batches\n#\n#     full_pred, pred, y, image_after_mask = get_image_for_plot(fake, img_with_erased_region, mask, original_img,\n#                                                               original_img)\n#\n#     print(f\"Validate Error:\\nAvg loss: {validate_loss:>8f}\\n\")\n#     return validate_loss, full_pred, y, pred, image_after_mask\n\n\ndef train_ae_only(dataloader, auto_encoder_pack):\n    auto_encoder, ae_criterion, ae_optimizer, a = auto_encoder_pack['model'], auto_encoder_pack['loss'], \\\n                                                  auto_encoder_pack[\n                                                      'optimizer'], auto_encoder_pack['lambda']\n    size = len(train_loader.batch_sampler)\n    pbar = tqdm(range(size))\n    auto_encoder.train()\n    train_ae_loss = 0\n\n    for batch, X in zip(pbar, dataloader):\n        original_img, erased_region, img_with_erased_region = X\n        X = img_with_erased_region.to(device)\n        real = original_img.to(device)\n\n        # Compute prediction error\n        auto_encoder.zero_grad()\n        fake = auto_encoder(X)\n\n        loss = ae_criterion(fake, real)\n\n        train_ae_loss += loss.item()\n\n        # Backpropagation\n        ae_optimizer.zero_grad()\n        loss.backward()\n        a.backward()\n        ae_optimizer.step()\n        #         if batch % 40 == 0:\n        #             loss, current = loss.item(), batch * len(X)\n        #             print(f\"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]\")\n        pbar.set_description(f\"Batch {batch} | AutoEncoder loss = {train_ae_loss / (batch + 1):>7f}\")\n\n    train_ae_loss /= len(dataloader)\n\n    index = random.randint(0, real.shape[0] - 1)\n    region = crop_image(fake)\n    pred = add_missing_region_to_image(img_with_erased_region[index], region[index])\n    pred = pred.permute(1, 2, 0).detach().cpu()\n    y = original_img[index].permute(1, 2, 0).cpu()\n\n    # print(f\"Train Error: \\nAvg auto encoder loss: {train_ae_loss:>8f}\")\n    return train_ae_loss, pred, y\n\n","metadata":{"id":"ZLlgt-rk0J-I","execution":{"iopub.status.busy":"2022-03-20T16:51:31.991459Z","iopub.execute_input":"2022-03-20T16:51:31.992539Z","iopub.status.idle":"2022-03-20T16:51:32.044882Z","shell.execute_reply.started":"2022-03-20T16:51:31.992500Z","shell.execute_reply":"2022-03-20T16:51:32.044031Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Model parameters","metadata":{"id":"NLFRxLNc0J-J"}},{"cell_type":"code","source":"def auto_encoder_parameters():\n\n    encoder = nn.Sequential(\n        nn.Conv2d(in_channels=3, out_channels=64, kernel_size=(4, 4), stride=(2, 2), padding=1),  # output: 64,128,128\n        nn.BatchNorm2d(64),\n        nn.PReLU(),\n        nn.Conv2d(in_channels=64, out_channels=128, kernel_size=(4, 4), stride=(2, 2), padding=1),  # output: 128,64,64\n        nn.BatchNorm2d(128),\n        nn.PReLU(),\n        nn.Conv2d(in_channels=128, out_channels=256, kernel_size=(4, 4), stride=(2, 2), padding=1),  # output: 256,32,32\n        nn.BatchNorm2d(256),\n        nn.PReLU(),\n        nn.Conv2d(in_channels=256, out_channels=512, kernel_size=(4, 4), stride=(2, 2), padding=1),  # output: 512,16,16\n        nn.BatchNorm2d(512),\n        nn.PReLU(),\n        nn.Conv2d(in_channels=512, out_channels=1024, kernel_size=(4, 4), stride=(2, 2), padding=1),  # output: 1024,8,8\n        nn.BatchNorm2d(1024),\n        nn.PReLU(),\n        nn.Conv2d(in_channels=1024, out_channels=2048, kernel_size=(4, 4), stride=(2, 2), padding=1),  # output: 2048,4,4\n        nn.BatchNorm2d(2048),\n        nn.PReLU(),\n        nn.Conv2d(in_channels=2048, out_channels=8192, kernel_size=(4, 4), stride=(1, 1), padding=0,groups=2048),  # output: 8192,1,1\n        nn.BatchNorm2d(8192),\n        nn.PReLU(),\n    )\n\n    # encoder =  torch.hub.load('pytorch/vision:v0.10.0', 'mobilenet_v2', pretrained=True)\n    decoder = nn.Sequential(\n        nn.ConvTranspose2d(in_channels=8192, out_channels=2048, kernel_size=(4, 4), stride=(2, 2)),  # output: 1024,4,4\n        nn.BatchNorm2d(2048),\n        nn.PReLU(),\n        nn.ConvTranspose2d(in_channels=2048, out_channels=1024, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1)),\n        # output: 512,8,8\n        nn.BatchNorm2d(1024),\n        nn.PReLU(),\n        nn.ConvTranspose2d(in_channels=1024, out_channels=512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1)),\n        # output: 512,8,8\n        nn.BatchNorm2d(512),\n        nn.PReLU(),\n        nn.ConvTranspose2d(in_channels=512, out_channels=256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1)),\n        # output: 256,16,16\n        nn.BatchNorm2d(256),\n        nn.PReLU(),\n        nn.ConvTranspose2d(in_channels=256, out_channels=128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1)),\n        # output: 128,32,32\n        nn.BatchNorm2d(128),\n        nn.PReLU(),\n        nn.ConvTranspose2d(in_channels=128, out_channels=64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1)),\n        # output: 64,64,64\n        nn.BatchNorm2d(64),\n        nn.PReLU(),\n        nn.ConvTranspose2d(in_channels=64, out_channels=3, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1)),\n        # output: 32,128,128\n        # nn.BatchNorm2d(32),\n        # nn.PReLU(),\n        # nn.ConvTranspose2d(in_channels=32, out_channels=3, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1)),  # output: 3,64,64\n        nn.Tanh(),\n        # nn.BatchNorm2d(3),\n    )\n    # decoder.apply(weights_init)\n    return decoder, encoder\n\n\ndef discriminator_parameters():\n    discriminator = nn.Sequential(\n        nn.Conv2d(in_channels=3, out_channels=16, kernel_size=(3, 3), stride=(1, 1), padding=1),\n        nn.BatchNorm2d(16),\n        nn.PReLU(),\n        nn.MaxPool2d((2, 2)),\n        nn.Conv2d(in_channels=16, out_channels=32, kernel_size=(3, 3), stride=(1, 1), padding=1),\n        nn.BatchNorm2d(32),\n        nn.PReLU(),\n        nn.MaxPool2d((2, 2)),\n        nn.Conv2d(in_channels=32, out_channels=64, kernel_size=(3, 3), stride=(1, 1), padding=1),\n        nn.BatchNorm2d(64),\n        nn.PReLU(),\n        nn.MaxPool2d((2, 2)),\n        nn.Conv2d(in_channels=64, out_channels=128, kernel_size=(3, 3), stride=(1, 1), padding=1),\n        nn.BatchNorm2d(128),\n        nn.PReLU(),\n        nn.MaxPool2d((2, 2)),\n        nn.Conv2d(in_channels=128, out_channels=256, kernel_size=(3, 3), stride=(1, 1), padding=1),\n        nn.BatchNorm2d(256),\n        nn.PReLU(),\n        nn.MaxPool2d((2, 2)),\n        nn.Conv2d(in_channels=256, out_channels=512, kernel_size=(3, 3), stride=(1, 1), padding=1),\n        nn.BatchNorm2d(512),\n        nn.PReLU(),\n        nn.MaxPool2d((2, 2)),\n        nn.Conv2d(in_channels=512, out_channels=1024, kernel_size=(3, 3), stride=(1, 1), padding=1),\n        nn.BatchNorm2d(1024),\n        nn.PReLU(),\n        nn.MaxPool2d((4, 4)),\n        nn.Flatten(),\n        nn.Linear(in_features=1024,out_features=1),\n        # nn.Conv2d(in_channels=1024, out_channels=1, kernel_size=(2, 2), stride=(1, 1)),\n        nn.Sigmoid(),\n\n    )\n    return discriminator","metadata":{"id":"u6eI0f1D0J-J","execution":{"iopub.status.busy":"2022-03-20T16:51:32.048065Z","iopub.execute_input":"2022-03-20T16:51:32.048286Z","iopub.status.idle":"2022-03-20T16:51:32.077006Z","shell.execute_reply.started":"2022-03-20T16:51:32.048259Z","shell.execute_reply":"2022-03-20T16:51:32.076154Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Main loop","metadata":{"id":"sUEWjYV20J-K"}},{"cell_type":"code","source":"def auto_encoder_only_loop(auto_encoder_pack, train_loader, validate_loader, epochs=10):\n    for t in range(epochs):\n        print(f\"Epoch {t + 1}\\n-------------------------------\")\n        loss, train_pred, train_y = train_ae_only(dataloader=train_loader, auto_encoder_pack=auto_encoder_pack)\n        loss, val_pred, val_y = validate(dataloader=validate_loader, model=auto_encoder_pack['model'],\n                                         loss_fn=nn.MSELoss())\n        show_images(train_y, train_pred, val_y, val_pred)\n\n\ndef main_loop(auto_encoder_pack, discriminator_pack, train_loader, validate_loader, scaler, epochs=50):\n    # load_weights(discriminator_pack['model'], 'd_weights.bin')\n    best_psnr = 0\n    best_mse = 1e6\n    train_loss, val_mse_loss, val_psnr = [], [] , []\n    for t in range(epochs):\n        print(f\"Epoch {t + 1}\\n-------------------------------\")\n        if scaler is None:\n            loss, train_full_pred, train_y, train_pred, train_image_after_mask = train(dataloader=train_loader,\n                                                                                       auto_encoder_pack=auto_encoder_pack,\n                                                                                       discriminator_pack=discriminator_pack)\n        else:\n            loss, train_full_pred, train_y, train_pred, train_image_after_mask = train_with_cuda(\n                dataloader=train_loader,\n                auto_encoder_pack=auto_encoder_pack,\n                discriminator_pack=discriminator_pack,\n                scaler=scaler)\n\n        train_loss.append(loss)\n        mse_loss,psnr, val_full_pred, val_y, val_pred, validate_image_after_mask = validate(dataloader=validate_loader,\n                                                                                   model=auto_encoder_pack['model'],\n                                                                                   loss_fn=nn.MSELoss())\n        val_mse_loss.append(mse_loss)\n        val_psnr.append(psnr)\n        show_with_missing_regions(train_full_pred, train_y, train_pred, train_image_after_mask, val_full_pred, val_y,\n                                  val_pred, validate_image_after_mask)\n\n        if best_psnr <  psnr:\n          best_psnr = psnr\n          save_models(auto_encoder_pack,discriminator_pack,loss,t,'psnr')\n        if best_mse > mse_loss:\n          best_mse = mse_loss\n          save_models(auto_encoder_pack,discriminator_pack,loss,t,'mse')\n\n\n\n    print(\"Done!\")\n","metadata":{"id":"f8H7sOTN0J-K","execution":{"iopub.status.busy":"2022-03-20T16:51:32.080689Z","iopub.execute_input":"2022-03-20T16:51:32.080928Z","iopub.status.idle":"2022-03-20T16:51:32.093899Z","shell.execute_reply.started":"2022-03-20T16:51:32.080903Z","shell.execute_reply":"2022-03-20T16:51:32.093122Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Identity(nn.Module):\n    def __init__(self):\n        super().__init__()\n        \n    def forward(self, x):\n        return x\n\nclass EncoderBlock(nn.Module):\n    def __init__(self, in_ch, out_ch, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), groups=1):\n        super().__init__()\n        self.conv = nn.Conv2d(in_channels=in_ch, out_channels=out_ch, kernel_size=kernel_size, stride=stride,\n                              padding=padding, groups=groups)\n        self.batch_norm = nn.BatchNorm2d(out_ch)\n        self.p_relu = nn.PReLU()\n\n    def forward(self, x):\n        return self.p_relu(self.batch_norm(self.conv(x)))\n\n\nclass DecoderBlock(nn.Module):\n    def __init__(self, in_ch, out_ch, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), output=False, batch_norm=True):\n        super().__init__()\n        self.conv = nn.ConvTranspose2d(in_channels=in_ch, out_channels=out_ch, kernel_size=kernel_size, stride=stride,\n                                       padding=padding)\n        self.batch_norm = nn.BatchNorm2d(out_ch)\n        self.conv2 = nn.Conv2d(in_channels=out_ch, out_channels=out_ch, kernel_size=1,stride=1,padding=0)\n        self.activition = nn.Tanh() if output else nn.PReLU()\n        self.norm = batch_norm\n        \n\n    def forward(self, x):\n        if self.norm:\n            return self.activition(self.batch_norm(self.conv(x))) # * self.activition2(self.conv(x))\n        else:\n            return self.activition(self.conv(x)) # * self.activition2(self.conv(x)) \n\n\nclass Encoder(nn.Module):\n    def __init__(self, chs=(3, 64, 128, 256, 512, 1024, 2048, 8092)):\n        super().__init__()\n        self.enc_blocks = nn.ModuleList([EncoderBlock(chs[i], chs[i + 1]) for i in range(len(chs) - 2)])\n        self.enc_blocks.append(EncoderBlock(chs[-2], chs[-1], padding=0, stride=(1, 1), groups=chs[-2]))\n\n\n    def forward(self, x):\n        ftrs = []\n        for block in self.enc_blocks:\n            x = block(x)\n            ftrs.append(x)\n        return ftrs\n\n\nclass Decoder(nn.Module):\n    def __init__(self, chs=(1024, 512, 256, 128, 64)):\n        super().__init__()\n        self.dec_blocks = nn.ModuleList([DecoderBlock(chs[i][0], chs[i][1]) for i in range(1,len(chs))])\n        self.dec_blocks.insert(0, DecoderBlock(chs[0][0], chs[0][1], output=True, padding=(0, 0), batch_norm=False))\n        self.chs = chs\n\n    def forward(self,x, encoder_features):\n\n        x = self.dec_blocks[0](encoder_features[0])\n        for i in range(1,len(self.chs) - 1):\n            x = torch.cat([x, encoder_features[i]], dim=1)\n            x = self.dec_blocks[i](x)\n\n        return x\n \nclass UNet(nn.Module):\n    def __init__(self, enc_chs=(3, 64, 128, 256, 512, 1024, 2048, 4096), dec_chs=([4096,2048], [4096,1024], [2048,512], [1024,256], [512,128], [256,64], [128,3],[6 ,3]),\n                 num_class=1, retain_dim=False, out_sz=(256, 256)):\n        super().__init__()\n\n        self.encoder = Encoder(enc_chs)\n        self.decoder = Decoder(dec_chs)\n        self.head = nn.Conv2d(enc_chs[-1], 4096, 1)\n        self.retain_dim = retain_dim\n        self.out_sz = out_sz\n        self.name = MODEL\n\n\n    def forward(self, x):\n        enc_ftrs = self.encoder(x)\n        out = self.head(enc_ftrs[-1])\n        out = self.decoder(out, enc_ftrs[::-1])\n        if self.retain_dim:\n            out = F.interpolate(out, self.out_sz)\n        return out\n\n\n\nclass ContentLoss(nn.Module):\n    \"\"\"Constructs a content loss function based on the VGG19 network.\n    Using high-level feature mapping layers from the latter layers will focus more on the texture content of the image.\n    Paper reference list:\n        -`Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network <https://arxiv.org/pdf/1609.04802.pdf>` paper.\n        -`ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks                    <https://arxiv.org/pdf/1809.00219.pdf>` paper.\n        -`Perceptual Extreme Super Resolution Network with Receptive Field Block               <https://arxiv.org/pdf/2005.12597.pdf>` paper.\n     \"\"\"\n\n    def __init__(self,normelize=True) -> None:\n        super(ContentLoss, self).__init__()\n        # Load the VGG19 model trained on the ImageNet dataset.\n        vgg19 = models.vgg19(pretrained=True).eval()\n        # Extract the thirty-sixth layer output in the VGG19 model as the content loss.\n        self.feature_extractor = nn.Sequential(*list(vgg19.features.children())[:36])\n        # Freeze model parameters.\n        for parameters in self.feature_extractor.parameters():\n            parameters.requires_grad = False\n\n        # The preprocessing method of the input data. This is the VGG model preprocessing method of the ImageNet dataset.        \n        # if normelize:\n        #     self.mean = torch.Tensor([-0.03, -0.088, -0.188]).view(1, 3, 1, 1).to(device)\n        #     self.std =  torch.Tensor([0.458, 0.448, 0.450]).view(1, 3, 1, 1).to(device)\n        # else:\n        self.mean = torch.Tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).to(device)\n        self.std =  torch.Tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).to(device)\n\n        \n\n    def forward(self, generated_image: Tensor, original: Tensor) -> Tensor:\n        loss = 0\n        x1 = generated_image.sub(self.mean).div(self.std)\n        x2 = original.sub(self.mean).div(self.std)\n        with torch.no_grad():\n          for l in self.feature_extractor:\n              x1 = l(x1)\n              x2 = l(x2)\n              if isinstance(l, nn.Conv2d):\n                  loss += F.l1_loss(x1,x2)\n        # loss = F.l1_loss(self.feature_extractor(generated_image), self.feature_extractor(original))\n\n        return loss\n","metadata":{"id":"KIB1sjzl0J-L","execution":{"iopub.status.busy":"2022-03-20T16:51:32.095061Z","iopub.execute_input":"2022-03-20T16:51:32.095495Z","iopub.status.idle":"2022-03-20T16:51:32.133626Z","shell.execute_reply.started":"2022-03-20T16:51:32.095440Z","shell.execute_reply":"2022-03-20T16:51:32.132844Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Get Data\ntrain_loader, validate_loader = get_data_from_files('../input/gan-getting-started')\n# show_sample(data_set)\n\n# create Discriminator Pack\ndiscriminator = Discriminator().to(device)\ndiscriminator.apply(weights_init)\noptimizerD = torch.optim.Adam(discriminator.parameters(), lr=0.00002)\n# discriminator_pack = {'model': discriminator, 'loss': nn.BCELoss(), 'optimizer': optimizerD}\ndiscriminator_pack = {'model': discriminator, 'loss': nn.BCEWithLogitsLoss(), 'optimizer': optimizerD}\n# create Auto Encoder Pack\nauto_encoder = AutoEncoder().to(device)\nauto_encoder.apply(weights_init)\nloss_lambda = nn.Parameter(torch.tensor(0.9985))\noptimizerAE = torch.optim.Adam(list(auto_encoder.parameters()), lr=0.0003)\ncontent_loss = ContentLoss().to(device) if CONTENT_LOSS else None\nauto_encoder_pack = {'model': auto_encoder, 'loss': nn.MSELoss(), 'optimizer': optimizerAE, 'lambda': loss_lambda,'content_loss':content_loss}\n\nscaler = None if device == 'cpu' else amp.GradScaler()\n# Run model\n# auto_encoder_only_loop(auto_encoder_pack, train_loader, validate_loader, epochs=10)\nmain_loop(auto_encoder_pack, discriminator_pack, train_loader, validate_loader,scaler,epochs = 100)","metadata":{"id":"fohl0lrO0J-K","outputId":"0612a8d4-e0ad-4e26-8d54-b0be9bb49a44","execution":{"iopub.status.busy":"2022-03-20T16:51:32.136949Z","iopub.execute_input":"2022-03-20T16:51:32.137885Z","iopub.status.idle":"2022-03-20T21:16:13.208418Z","shell.execute_reply.started":"2022-03-20T16:51:32.137832Z","shell.execute_reply":"2022-03-20T21:16:13.207707Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"main_loop(auto_encoder_pack, discriminator_pack, train_loader, validate_loader,scaler,epochs = 100)","metadata":{"id":"M1ronSSLS8qG","execution":{"iopub.status.busy":"2022-03-20T21:16:13.212514Z","iopub.execute_input":"2022-03-20T21:16:13.213813Z","iopub.status.idle":"2022-03-20T23:34:38.433178Z","shell.execute_reply.started":"2022-03-20T21:16:13.213768Z","shell.execute_reply":"2022-03-20T23:34:38.431223Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def test(dataloader,model):\n\n    def show_sample(source_img,mask,fake):\n        source_img = source_img.detach().cpu()\n        mask = mask.detach().cpu()\n        fake = fake.detach().cpu()\n        full_pred = torch.where(mask == 1, fake, source_img)\n        pred = torch.where(mask == 0, torch.ones_like(full_pred), fake)\n        pred = torch.where(mask == 1, fake, pred)\n        full_pred = full_pred.permute(1, 2, 0).detach().cpu()\n        pred = pred.permute(1, 2, 0).detach().cpu()\n        source_img = source_img.permute(1,2,0)\n        fig, axs = plt.subplots(1, 3)\n        fig.set_dpi(120)\n        axs[0].imshow(source_img)\n        axs[1].imshow(pred)\n        axs[2].imshow(full_pred)\n        for x in axs:\n            x.axis('off')\n        plt.subplots_adjust(hspace=0, wspace=0)\n        plt.tight_layout()\n        plt.show()\n\n    model.eval()\n    with torch.no_grad():\n        for X in dataloader:\n            source_img, mask = X\n            X = source_img.to(device)\n            fake = model(X)\n            show_sample(source_img[0],mask[0],fake[0])\n","metadata":{"execution":{"iopub.status.busy":"2022-03-20T23:34:51.968857Z","iopub.execute_input":"2022-03-20T23:34:51.969130Z","iopub.status.idle":"2022-03-20T23:34:51.983085Z","shell.execute_reply.started":"2022-03-20T23:34:51.969101Z","shell.execute_reply":"2022-03-20T23:34:51.981796Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"checkpoint = torch.load('./auto_encoder_center_psnr')\n\n# discriminator = Discriminator().to(device)\n# optimizerD = torch.optim.Adam(discriminator.parameters(), lr=0.00002)\nauto_encoder = AutoEncoder().to(device)\n# optimizerAE = torch.optim.Adam(list(auto_encoder.parameters()), lr=0.0003)\n\n# epoch = checkpoint['epoch']\n# loss = checkpoint['loss']\n\n# discriminator.load_state_dict(checkpoint['d_state_dict'])\nauto_encoder.load_state_dict(checkpoint['ae_state_dict'])\n# optimizerD.load_state_dict(checkpoint['d_optimizer_state_dict'])\n# optimizerAE.load_state_dict(checkpoint['ae_optimizer_state_dict'])\n","metadata":{"id":"MRLxEBw3S9JV","execution":{"iopub.status.busy":"2022-03-20T23:34:38.436664Z","iopub.status.idle":"2022-03-20T23:34:38.437416Z","shell.execute_reply.started":"2022-03-20T23:34:38.437167Z","shell.execute_reply":"2022-03-20T23:34:38.437194Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_dataset = ImageTestDataSet('../input/deep-learning-project/test/test', masks_dir=MASK_DIR,transforms=general_transformer())\nprint(len(test_dataset))\ntest_loader = DataLoader(test_dataset, batch_size=1)\ntest(test_loader,auto_encoder)","metadata":{"id":"04xBrH-mS9au","execution":{"iopub.status.busy":"2022-03-20T23:35:02.384090Z","iopub.execute_input":"2022-03-20T23:35:02.384348Z","iopub.status.idle":"2022-03-20T23:35:02.395404Z","shell.execute_reply.started":"2022-03-20T23:35:02.384319Z","shell.execute_reply":"2022-03-20T23:35:02.394388Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"torch.cuda.empty_cache()\nimport gc\ngc.collect()\n","metadata":{"execution":{"iopub.status.busy":"2022-03-20T23:34:38.440395Z","iopub.status.idle":"2022-03-20T23:34:38.441243Z","shell.execute_reply.started":"2022-03-20T23:34:38.440986Z","shell.execute_reply":"2022-03-20T23:34:38.441012Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<a href=\"./auto_encoder_center_psnr\"> Download File </a>","metadata":{}},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}