{"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 os\n# make folder for the outputs of the model\nos.mkdir(\"/kaggle/working/images\")\nos.mkdir(\"./validation\")\n","metadata":{"id":"QHnfOMqY6Ctg","executionInfo":{"status":"ok","timestamp":1625143463808,"user_tz":-180,"elapsed":24387,"user":{"displayName":"Eitan Kerzhner","photoUrl":"https://lh3.googleusercontent.com/a-/AOh14GhQc5XyxIsXL6f1XLggbRylGSlAVO4fkoAcoAUPog=s64","userId":"10997308535591103312"}},"outputId":"9d5eba5a-202d-4c24-b3ab-9546d7acb944","execution":{"iopub.status.busy":"2021-07-31T14:52:05.220885Z","iopub.execute_input":"2021-07-31T14:52:05.221257Z","iopub.status.idle":"2021-07-31T14:52:05.232451Z","shell.execute_reply.started":"2021-07-31T14:52:05.221179Z","shell.execute_reply":"2021-07-31T14:52:05.231479Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Configuration","metadata":{}},{"cell_type":"code","source":"import torch\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nBATCH_SIZE = 8\nLEARNING_RATE = 0.00021\nLAMBDA_IDENTITY = 0.1\nLAMBDA_CYCLE = 15\nNUM_WORKERS = 4\nNUM_EPOCHS = 200\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","metadata":{"id":"KuYh_2CX6J-H","execution":{"iopub.status.busy":"2021-07-31T14:52:05.234037Z","iopub.execute_input":"2021-07-31T14:52:05.234411Z","iopub.status.idle":"2021-07-31T14:52:08.026209Z","shell.execute_reply.started":"2021-07-31T14:52:05.234377Z","shell.execute_reply":"2021-07-31T14:52:08.025295Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Objects for the model:\n1. Generator: We create a generetor with convolotion blockes that shrink the image dimentions to shring representation of the image, and then expand them to the other domain image. \n2. Discriminator: We create a discriminator with convolotion blocks, that give us a representian of the image in low dimentions, and then we decide is the image is origin or fake with sigmoid on the shring representation. This way the discriminator learn to represente the image according to the domain that he need to decide if it real or fake. ","metadata":{}},{"cell_type":"code","source":"import torch.nn as nn\nimport torchvision.models as models\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","metadata":{"id":"ldAjd4Q06N5N","execution":{"iopub.status.busy":"2021-07-31T14:52:08.029662Z","iopub.execute_input":"2021-07-31T14:52:08.029945Z","iopub.status.idle":"2021-07-31T14:52:08.045163Z","shell.execute_reply.started":"2021-07-31T14:52:08.029918Z","shell.execute_reply":"2021-07-31T14:52:08.044355Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch.nn as nn\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","metadata":{"id":"dg4SnwYW6RyY","execution":{"iopub.status.busy":"2021-07-31T14:52:08.046776Z","iopub.execute_input":"2021-07-31T14:52:08.047389Z","iopub.status.idle":"2021-07-31T14:52:08.059542Z","shell.execute_reply.started":"2021-07-31T14:52:08.047348Z","shell.execute_reply":"2021-07-31T14:52:08.058753Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"1. Transforms: we create transform that will generate as different photo in each call, the transform is probablistic and have caple of different options for output. This way we can get good perfomence even when we use only 30 photo in our data set.\n2. MonetRealDataset: We create a data set for this project, the data set provide us each time one photo of monet, and one photo from the real photos, and the names of the photos.","metadata":{}},{"cell_type":"code","source":"from torch.utils.data import Dataset\nimport numpy as np\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 =False):\n        # monet = monet\n        # real == real\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)) # 1000, 1500\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","metadata":{"id":"Ccly6FRKFAZA","execution":{"iopub.status.busy":"2021-07-31T14:52:08.060816Z","iopub.execute_input":"2021-07-31T14:52:08.061143Z","iopub.status.idle":"2021-07-31T14:52:08.078026Z","shell.execute_reply.started":"2021-07-31T14:52:08.061109Z","shell.execute_reply":"2021-07-31T14:52:08.077218Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"* Data_loading: We create a data set for train and for validation, and create the data loader for each one.\n* Initialize: The function initialize all the models parts - discriminators and generators, and the loss function that we will use in the model.\n* Dave_checkpoint:The function save us the model for future use.\n* Load_checkpoint:The function load us the model and return as model that we can use.","metadata":{}},{"cell_type":"code","source":"def 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\n\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    # initialize opt gen\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    # initialize loss\n    L1 = nn.L1Loss()\n    mse = nn.MSELoss()\n    \n    return disc_H, disc_Z, gen_Z, gen_H, opt_disc, opt_gen, L1, mse\n\n\ndef save_checkpoint(model, optimizer, filename=\"my_checkpoint.pth.tar\"):\n    print(\"=> Saving checkpoint\")\n    checkpoint = {\n        \"state_dict\": model.state_dict(),\n        \"optimizer\": optimizer.state_dict(),\n    }\n    torch.save(checkpoint, filename)\n\n\ndef load_checkpoint(checkpoint_file, model, optimizer, lr):\n    print(\"=> Loading checkpoint\")\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    # If we don't do this then it will just have learning rate of old checkpoint\n    # and it will lead to many hours of debugging \\:\n    for param_group in optimizer.param_groups:\n        param_group[\"lr\"] = lr\n","metadata":{"id":"D0d-xrY4Fc1v","execution":{"iopub.status.busy":"2021-07-31T14:52:08.079412Z","iopub.execute_input":"2021-07-31T14:52:08.079842Z","iopub.status.idle":"2021-07-31T14:52:08.093669Z","shell.execute_reply.started":"2021-07-31T14:52:08.079806Z","shell.execute_reply":"2021-07-31T14:52:08.09281Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"* Prediction: The function get the model and the data loader, and for each photo in the data loader, it create the fake monet of the photo and save it. In the end of the function we create a zip file with all the photos that we made from the model.\n* Train: The function get the model and the data loader, and run one epoch of the data loader on the model. We first training the generators, and after this we training the disciminators, and in the end we complite the cycle by taking the fake photo and use the other generetor to return it to the origin photo, and we use min squre error as a loss function for the cycle.\n* End_of_epoch: At the end of the epoch, we chosing 4 photo and use the generetor with those photo, and present it to the user to see in the eyes the model output in the end of the epoch.","metadata":{}},{"cell_type":"code","source":"import matplotlib.pyplot as plt\nimport matplotlib.image as mpimg\nimport shutil\nimport PIL\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        # Train Discriminators H and Z\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    # make zip for submit\n    shutil.make_archive(\"/kaggle/working/images\", 'zip', OUTPUT_PREDICT_PATH)\n\nGlobalIdx =0\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        # Train Discriminators H and Z\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            # put it togethor\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        # Train Generators H and Z\n        with torch.cuda.amp.autocast():\n            # adversarial loss for both generators\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 loss\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 loss (remove these for efficiency if you set lambda_identity=0)\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            # add all togethor\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_checkpoint(gen_H, opt_gen, filename=CHECKPOINT_GEN_H)\n                save_checkpoint(gen_Z, opt_gen, filename=CHECKPOINT_GEN_Z)\n                save_checkpoint(disc_H, opt_disc, filename=CHECKPOINT_CRITIC_H)\n                save_checkpoint(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(4):\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 ","metadata":{"id":"N4-yaij_6Ugg","execution":{"iopub.status.busy":"2021-07-31T14:52:08.094867Z","iopub.execute_input":"2021-07-31T14:52:08.095217Z","iopub.status.idle":"2021-07-31T14:52:08.120824Z","shell.execute_reply.started":"2021-07-31T14:52:08.095165Z","shell.execute_reply":"2021-07-31T14:52:08.119346Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import sys\nimport torch\n#import tensorflow.Tensor as Tensor\nfrom fastai.vision.all import show_image\n#import torch.Tensor as Tensor\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\nfrom PIL import Image\nimport os\nfrom torch.utils.data import Dataset\nimport numpy as np\n#from IPython.display import display, Image as display, ImageD\n\ndef main():\n\n    # initialize\n    disc_H, disc_Z, gen_Z, gen_H, opt_disc, opt_gen, L1, mse  = initialize()\n    \n    # loading params\n    if LOAD_MODEL:\n        load_checkpoint(\n           CHECKPOINT_GEN_H, gen_H, opt_gen, LEARNING_RATE,\n        )\n        load_checkpoint(\n            CHECKPOINT_GEN_Z, gen_Z, opt_gen, LEARNING_RATE,\n        )\n        load_checkpoint(\n            CHECKPOINT_CRITIC_H, disc_H, opt_disc, LEARNING_RATE,\n        )\n        load_checkpoint(\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    # start training\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        for i in range(len(total_loss)):\n            iter.append(i)\n        plt.plot(iter, total_loss, color='black', linestyle='dashed', linewidth=3,\n                marker='o', markerfacecolor='gray', markersize=4)\n        plt.xlabel(\"Iteration Number\")\n        plt.ylabel(\"Loss\")\n        plt.title(\"Loss Rate\")\n        plt.show()\n    if PREDICTION:\n        prediction(gen_Z, val_loader)\n  \n\nmain()\n","metadata":{"id":"0OfmRKRw6cxK","outputId":"06789c2a-8665-4376-f10f-2f2ff27bb216","execution":{"iopub.status.busy":"2021-07-31T14:52:08.123115Z","iopub.execute_input":"2021-07-31T14:52:08.123474Z"},"trusted":true},"execution_count":null,"outputs":[]}]}