{"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.nn as nn\nimport torchvision.models as models\nimport torch\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport os\nfrom torch.utils.data import Dataset\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport matplotlib.image as mpimg\nimport shutil\nimport PIL\nfrom PIL import Image\n\nimport sys\nfrom torch.utils.data import DataLoader\nimport torch.nn as nn\nimport torch.optim as optim\nfrom tqdm import tqdm\nfrom torchvision.utils import save_image\n\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nos.mkdir(\"/kaggle/working/images\")\nos.mkdir(\"./validation\")\nBATCH_SIZE = 8\nLEARNING_RATE = 0.00021\nLAMBDA_IDENTITY = 0.1\nLAMBDA_CYCLE = 15\nNUM_WORKERS = 4\nNUM_EPOCHS = 46\nLOAD_MODEL = False\nSAVE_MODEL = False\nPREDICTION = True\nCHECKPOINT_GEN_H = \"../input/model-state/genh.pth.tar\"\nCHECKPOINT_GEN_Z = \"../input/model-state/genz.pth.tar\"\nCHECKPOINT_CRITIC_H = \"../input/model-state/critich.pth.tar\"\nCHECKPOINT_CRITIC_Z = \"../input/model-state/criticz.pth.tar\"\nTRIAN_PATH = \"../input/gan-getting-started\"\nVALID_PATH = \"../input/gan-getting-started\"\nOUTPUT_TRAIN_PATH=\"./train\"\nOUTPUT_PREDICT_PATH=\"./validation\"\ntorch.cuda.empty_cache()\n\n\nclass ConvBlock(nn.Module):\n    def __init__(self, in_channels, out_channels, down=True, use_act=True, **kwargs):\n        super().__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, padding_mode=\"reflect\", **kwargs)\n            if down\n            else nn.ConvTranspose2d(in_channels, out_channels, **kwargs),\n            nn.InstanceNorm2d(out_channels),\n            nn.ReLU(inplace=True) if use_act else nn.Identity()\n        )\n\n    def forward(self, x):\n        return self.conv(x)\n\n\nclass ResidualBlock(nn.Module):\n    def __init__(self, channels):\n        super().__init__()\n        self.block = nn.Sequential(\n            ConvBlock(channels, channels, kernel_size=3, padding=1),\n            ConvBlock(channels, channels, use_act=False, kernel_size=3, padding=1),\n        )\n\n    def forward(self, x):\n        return x + self.block(x)\n\nclass Generator(nn.Module):\n    def __init__(self, img_channels, num_features = 64, num_residuals=9):\n        super().__init__()\n        self.initial = nn.Sequential(\n            nn.Conv2d(img_channels, num_features, kernel_size=7, stride=1, padding=3, padding_mode=\"reflect\"),\n            nn.InstanceNorm2d(num_features),\n            nn.ReLU(inplace=True),\n        )\n        self.down_blocks = nn.ModuleList(\n            [\n                ConvBlock(num_features, num_features*2, kernel_size=3, stride=2, padding=1),\n                ConvBlock(num_features*2, num_features*4, kernel_size=3, stride=2, padding=1),\n            ]\n        )\n        self.res_blocks = nn.Sequential(\n            *[ResidualBlock(num_features*4) for _ in range(num_residuals)]\n        )\n        self.up_blocks = nn.ModuleList(\n            [\n                ConvBlock(num_features*4, num_features*2, down=False, kernel_size=3, stride=2, padding=1, output_padding=1),\n                ConvBlock(num_features*2, num_features*1, down=False, kernel_size=3, stride=2, padding=1, output_padding=1),\n            ]\n        )\n\n        self.last = nn.Conv2d(num_features*1, img_channels, kernel_size=7, stride=1, padding=3, padding_mode=\"reflect\")\n\n    def forward(self, x):\n        x = self.initial(x)\n        for layer in self.down_blocks:\n            x = layer(x)\n        x = self.res_blocks(x)\n        for layer in self.up_blocks:\n            x = layer(x)\n        return torch.tanh(self.last(x))\n\nclass Block(nn.Module):\n    def __init__(self, in_channels, out_channels, stride):\n        super().__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, 4, stride, 1, bias=True, padding_mode=\"reflect\"),\n            nn.InstanceNorm2d(out_channels),\n            nn.LeakyReLU(0.2, inplace=True),\n        )\n\n    def forward(self, x):\n        return self.conv(x)\n\n\nclass Discriminator(nn.Module):\n    def __init__(self, in_channels=3, features=[64, 128, 256, 512]):\n        super().__init__()\n        self.initial = nn.Sequential(\n            nn.Conv2d(\n                in_channels,\n                features[0],\n                kernel_size=4,\n                stride=2,\n                padding=1,\n                padding_mode=\"reflect\",\n            ),\n            nn.LeakyReLU(0.2, inplace=True),\n        )\n\n        layers = []\n        in_channels = features[0]\n        for feature in features[1:]:\n            layers.append(Block(in_channels, feature, stride=1 if feature==features[-1] else 2))\n            in_channels = feature\n        layers.append(nn.Conv2d(in_channels, 1, kernel_size=4, stride=1, padding=1, padding_mode=\"reflect\"))\n        self.model = nn.Sequential(*layers)\n\n    def forward(self, x):\n        x = self.initial(x)\n        return torch.sigmoid(self.model(x))\n\n\ntrianTransforms = A.Compose(\n    [\n        A.Resize(width=256, height=256),\n        A.HorizontalFlip(p=0.5),\n        A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], max_pixel_value=255),\n        ToTensorV2(),\n     ],\n    additional_targets={\"image0\": \"image\"},\n)\npredictTransforms = A.Compose(\n    [\n        A.Resize(width=256, height=256),\n        A.OneOf([\n            A.Flip(p=1),\n            A.OpticalDistortion(p=1),\n            A.GlassBlur(p=1)\n        ], p=0.35),\n        A.OneOf([\n            A.GaussianBlur(p=1),\n            A.FancyPCA(p=1),\n        ], p=0.35),\n        A.HorizontalFlip(p=0.5),\n        A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], max_pixel_value=255),\n        ToTensorV2(),\n     ],\n    additional_targets={\"image0\": \"image\"},\n)\n\nclass MonetRealDataset(Dataset):\n    def __init__(self, root_monet, root_real, transform=None,limit =True):\n        self.root_monet = root_monet\n        self.root_real = root_real\n        self.transform = transform\n\n        self.monet_images = os.listdir(root_monet)\n        if limit:\n            self.monet_images = self.monet_images[:30]\n        self.real_images = os.listdir(root_real)\n        self.length_dataset = max(len(self.monet_images), len(self.real_images))\n        self.monet_len = len(self.monet_images)\n        self.real_len = len(self.real_images)\n\n    def __len__(self):\n        return self.length_dataset\n\n    def __getitem__(self, index):\n        monet_name = self.monet_images[index % self.monet_len]\n        real_name = self.real_images[index % self.real_len]\n\n        monet_path = os.path.join(self.root_monet, monet_name)\n        real_path = os.path.join(self.root_real, real_name)\n\n        monet_img = np.array(Image.open(monet_path).convert(\"RGB\"))\n        real_img = np.array(Image.open(real_path).convert(\"RGB\"))\n\n        if self.transform:\n            augmentations = self.transform(image=monet_img, image0=real_img)\n            monet_img = augmentations[\"image\"]\n            real_img = augmentations[\"image0\"]\n\n        return real_img, monet_img,real_name, monet_name\n\ndef data_loading():\n\n    dataset = MonetRealDataset(\n        root_real= TRIAN_PATH + \"/monet_jpg\", root_monet=TRIAN_PATH+\"/photo_jpg\", transform=predictTransforms, limit=True\n    )\n    val_dataset = MonetRealDataset(\n       root_real= VALID_PATH + \"/monet_jpg\", root_monet= VALID_PATH + \"/photo_jpg\", transform=trianTransforms\n    )\n    val_loader = DataLoader(\n        val_dataset,\n        batch_size=1,\n        shuffle=False,\n        pin_memory=True,\n    )\n    loader = DataLoader(\n        dataset,\n        batch_size=BATCH_SIZE,\n        shuffle=True,\n        num_workers=NUM_WORKERS,\n        pin_memory=True\n    )\n    \n\n    return dataset, val_dataset, val_loader, loader\n\nGlobalIdx =0\n\ndef initialize():\n    disc_H = Discriminator(in_channels=3).to(DEVICE)\n    disc_Z = Discriminator(in_channels=3).to(DEVICE)\n    gen_Z = Generator(img_channels=3, num_residuals=9).to(DEVICE)\n    gen_H = Generator(img_channels=3, num_residuals=9).to(DEVICE)\n    opt_disc = optim.Adam(\n        list(disc_H.parameters()) + list(disc_Z.parameters()),\n        lr=LEARNING_RATE,\n        betas=(0.5, 0.999),\n    )\n    opt_gen = optim.Adam(\n        list(gen_Z.parameters()) + list(gen_H.parameters()),\n        lr=LEARNING_RATE,\n        betas=(0.5, 0.999),\n    )\n    L1 = nn.L1Loss()\n    mse = nn.MSELoss()\n    return disc_H, disc_Z, gen_Z, gen_H, opt_disc, opt_gen, L1, mse\n\ndef save_state(model, optimizer, filename=\"my_checkpoint.pth.tar\"):\n    checkpoint = {\n        \"state_dict\": model.state_dict(),\n        \"optimizer\": optimizer.state_dict(),\n    }\n    torch.save(checkpoint, filename)\n\ndef load_state(checkpoint_file, model, optimizer, lr):\n    checkpoint = torch.load(checkpoint_file, map_location=DEVICE)\n    model.load_state_dict(checkpoint[\"state_dict\"])\n    optimizer.load_state_dict(checkpoint[\"optimizer\"])\n\n    for param_group in optimizer.param_groups:\n        param_group[\"lr\"] = lr\n\ndef prediction(gen_Z,loader):\n    H_reals = 0\n    H_fakes = 0\n    loop = tqdm(loader, leave=True)\n    for idx, (monet, real,monet_name,real_name) in enumerate(loop):\n        real = real.to(DEVICE)\n        with torch.cuda.amp.autocast():\n            fake_monet = gen_Z(real)\n            save_image(fake_monet*0.5+0.5, f\"{OUTPUT_PREDICT_PATH}/{idx}.png\")\n    shutil.make_archive(\"/kaggle/working/images\", 'zip', OUTPUT_PREDICT_PATH)\n\ndef train_fn(disc_H, disc_Z, gen_Z, gen_H, loader, opt_disc, opt_gen, l1, mse, d_scaler, g_scaler, num_epoch):\n    global GlobalIdx\n    H_reals = 0\n    H_fakes = 0\n    fake_real = None\n    fake_monet = None\n    loop = tqdm(loader, leave=True)\n\n    for idx, (monet, real, monet_name, real_name ) in enumerate(loop):\n        monet = monet.to(DEVICE)\n        real = real.to(DEVICE)\n\n        with torch.cuda.amp.autocast():\n            fake_real = gen_H(monet)\n            D_H_real = disc_H(real)\n            D_H_fake = disc_H(fake_real.detach())\n            H_reals += D_H_real.mean().item()\n            H_fakes += D_H_fake.mean().item()\n            D_H_real_loss = mse(D_H_real, torch.ones_like(D_H_real))\n            D_H_fake_loss = mse(D_H_fake, torch.zeros_like(D_H_fake))\n            D_H_loss = D_H_real_loss + D_H_fake_loss\n\n            fake_monet = gen_Z(real)\n            D_Z_real = disc_Z(monet)\n            D_Z_fake = disc_Z(fake_monet.detach())\n            D_Z_real_loss = mse(D_Z_real, torch.ones_like(D_Z_real))\n            D_Z_fake_loss = mse(D_Z_fake, torch.zeros_like(D_Z_fake))\n            D_Z_loss = D_Z_real_loss + D_Z_fake_loss\n\n            D_loss = (D_H_loss + D_Z_loss)/2\n\n        opt_disc.zero_grad()\n        d_scaler.scale(D_loss).backward()\n        d_scaler.step(opt_disc)\n        d_scaler.update()\n\n        with torch.cuda.amp.autocast():\n            D_H_fake = disc_H(fake_real)\n            D_Z_fake = disc_Z(fake_monet)\n            loss_G_H = mse(D_H_fake, torch.ones_like(D_H_fake))\n            loss_G_Z = mse(D_Z_fake, torch.ones_like(D_Z_fake))\n\n            cycle_monet = gen_Z(fake_real)\n            cycle_real = gen_H(fake_monet)\n            cycle_monet_loss = l1(monet, cycle_monet)\n            cycle_real_loss = l1(real, cycle_real)\n\n            identity_monet = gen_Z(monet)\n            identity_real = gen_H(real)\n            identity_monet_loss = l1(monet, identity_monet)\n            identity_real_loss = l1(real, identity_real)\n\n            G_loss = (\n                loss_G_Z\n                + loss_G_H\n                + cycle_monet_loss * LAMBDA_CYCLE\n                + cycle_real_loss * LAMBDA_CYCLE\n                + identity_real_loss * LAMBDA_IDENTITY\n                + identity_monet_loss * LAMBDA_IDENTITY\n            )\n\n        opt_gen.zero_grad()\n        g_scaler.scale(G_loss).backward()\n        g_scaler.step(opt_gen)\n        g_scaler.update()\n        if idx % 200 == 0:\n            if SAVE_MODEL and idx !=0:\n                save_state(gen_H, opt_gen, filename=CHECKPOINT_GEN_H)\n                save_state(gen_Z, opt_gen, filename=CHECKPOINT_GEN_Z)\n                save_state(disc_H, opt_disc, filename=CHECKPOINT_CRITIC_H)\n                save_state(disc_Z, opt_disc, filename=CHECKPOINT_CRITIC_Z)\n        loop.set_postfix(H_real=H_reals/(idx+1), H_fake=H_fakes/(idx+1))\n        GlobalIdx += 1\n    end_of_epoch(num_epoch,loader,gen_Z)\n    return G_loss\n    \ndef end_of_epoch(epoch,loader, gen_monet):\n    _, axel = plt.subplots(4, 2, figsize=(10, 15))\n    for i in range(2):\n        (monet, real, monet_name, real_name ) = next(iter(loader))\n        with torch.cuda.amp.autocast():\n            device_real = real.to(DEVICE)\n            prediction = gen_monet(device_real)[0]\n            prediction = prediction.cpu().detach().numpy()\n            prediction = (prediction * 127.5 + 127.5).astype(np.uint8)\n            prediction= np.moveaxis(prediction, 0, -1)\n            real = (real[0] * 127.5 + 127.5).numpy().astype(np.uint8)\n            real= np.moveaxis(real, 0, -1)\n            axel[i, 0].imshow(real)\n            axel[i, 1].imshow(prediction)\n            axel[i, 0].set_title(\"Input image\")\n            axel[i, 1].set_title(\"Monet\")\n            axel[i, 0].axis(\"off\")\n            axel[i, 1].axis(\"off\")\n    \n    plt.show()\n    plt.close()\n    \n\ndef main():\n    disc_H, disc_Z, gen_Z, gen_H, opt_disc, opt_gen, L1, mse  = initialize()\n    \n    if LOAD_MODEL:\n        load_state(\n           CHECKPOINT_GEN_H, gen_H, opt_gen, LEARNING_RATE,\n        )\n        load_state(\n            CHECKPOINT_GEN_Z, gen_Z, opt_gen, LEARNING_RATE,\n        )\n        load_state(\n            CHECKPOINT_CRITIC_H, disc_H, opt_disc, LEARNING_RATE,\n        )\n        load_state(\n            CHECKPOINT_CRITIC_Z, disc_Z, opt_disc, LEARNING_RATE,\n        )\n    dataset, val_dataset, val_loader, loader = data_loading()\n    \n    g_scaler = torch.cuda.amp.GradScaler()\n    d_scaler = torch.cuda.amp.GradScaler()\n    total_loss = []\n    for epoch in range(NUM_EPOCHS):\n        loss = train_fn(disc_H, disc_Z, gen_Z, gen_H, loader, opt_disc, opt_gen, L1, mse, d_scaler, g_scaler,epoch)\n        total_loss.append(loss)\n        iter = []\n    if PREDICTION:\n        prediction(gen_Z, val_loader)\n  \n\nif __name__ == \"__main__\":\n    main()","metadata":{"id":"6dVz1zlaXcLm","outputId":"37d62ac1-55ed-45c4-b526-49dbbcc34ad9","execution":{"iopub.status.busy":"2021-08-28T11:37:23.135271Z","iopub.execute_input":"2021-08-28T11:37:23.135649Z","iopub.status.idle":"2021-08-28T11:46:39.138769Z","shell.execute_reply.started":"2021-08-28T11:37:23.135551Z","shell.execute_reply":"2021-08-28T11:46:39.137818Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"id":"NVvW1HNIzWbH"},"execution_count":null,"outputs":[]}]}