{"cells":[{"metadata":{},"cell_type":"markdown","source":"**Inspired from - https://www.tensorflow.org/tutorials/generative/pix2pix**"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\nimport datetime\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session\n\nimport csv\nimport numpy as np\nimport os\nimport random\nimport time\nimport  seaborn as sns\nimport tifffile as tiff\nimport matplotlib\nmatplotlib.use('Agg')\nimport glob\nfrom matplotlib import pyplot as plt\nfrom PIL import Image\nimport cv2\nfrom skimage.transform import resize\nfrom skimage.morphology import binary_opening, disk, skeletonize\nfrom skimage.measure import label, regionprops\nimport pandas as pd\nfrom skimage.filters import threshold_otsu\ntry:\n    from skimage.io import imsave as imwrite\n    from skimage.io import imread\nexcept ImportError:\n    from cv2 import imread\n    from cv2 import imwrite\n\nfrom sklearn.model_selection import train_test_split","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install -q -U tensorboard","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def mask2rle(img):\n    '''\n    img: numpy array, 1 - mask, 0 - background\n    Returns run length as string formated\n    '''\n    pixels = img.T.flatten()\n    pixels = np.concatenate([[0], pixels, [0]])\n    runs = np.where(pixels[1:] != pixels[:-1])[0] + 1\n    runs[1::2] -= runs[::2]\n    return ' '.join(str(x) for x in runs)\n\n\ndef rle2mask(mask_rle, shape):\n    '''\n    mask_rle: run-length as string formated (start length)\n    shape: (width,height) of array to return\n    Returns numpy array, 1 - mask, 0 - background\n\n    '''\n    s = mask_rle.split()\n    starts, lengths = [np.asarray(x, dtype=int) for x in (s[0:][::2], s[1:][::2])]\n    starts -= 1\n    ends = starts + lengths\n    img = np.zeros(shape[0] * shape[1], dtype=np.uint8)\n    for lo, hi in zip(starts, ends):\n        img[lo:hi] = 1\n    return img.reshape(shape).T","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import csv\nimport numpy as np\nimport os\nimport random\nimport time\nimport  seaborn as sns\nimport tifffile as tiff\nimport matplotlib\nmatplotlib.use('Agg')\nimport glob\nfrom matplotlib import pyplot as plt\nfrom PIL import Image\nimport cv2\nfrom skimage.transform import resize\nfrom skimage.morphology import binary_opening, disk, skeletonize\nfrom skimage.measure import label, regionprops\nimport pandas as pd\nfrom skimage.filters import threshold_otsu\ntry:\n    from skimage.io import imsave as imwrite\n    from skimage.io import imread\nexcept ImportError:\n    from cv2 import imread\n    from cv2 import imwrite\n\nfrom sklearn.model_selection import train_test_split\n\nclass KDatasetPreparation():\n    def __init__(self,train_csv,\n                      train_img_folder,\n                      resize_factor = 1,\n                      point_sampling_level = 2,\n                      superImposed = True,\n                      patch_size = 512, #read_size\n                      stride = 512,\n                      save_gan_pair_list = True,\n                      save_patch = True,\n                      save_folder=None):\n        self.train_df = pd.read_csv(train_csv)\n        self.resize_factor = resize_factor\n        self.save_folder =save_folder\n        self.point_sampling_level =point_sampling_level\n        self.train_img_folder = train_img_folder\n        self.superImposed = superImposed\n        self.patch_size = patch_size\n        self.stride = stride\n        self.save_patch = save_patch\n        self.kernel = disk(2)\n        if self.save_folder is None:\n            self.save_folder = './tmp'\n        self.gan_pair_list = []\n        self.wsi_list = []\n        self.save_img_mask =  os.path.join(self.save_folder,'imgdata')\n        self.save_train_img =  os.path.join(self.save_folder,'patch')\n        self.save_train_mask =  os.path.join(self.save_folder,'mask')\n        os.makedirs(self.save_folder,exist_ok=True)\n        os.makedirs(self.save_img_mask,exist_ok=True)\n        os.makedirs(self.save_train_img,exist_ok=True)\n        os.makedirs(self.save_train_mask,exist_ok=True)\n        \n        self.save_gan_pair_list = save_gan_pair_list\n\n        print(self.train_df.head())\n\n\n    def get_mask_from_submission(self,subfilepath,testfolder,tryboth=True):\n        '''\n        test mask and rle\n        '''\n        sub_df = pd.read_csv(subfilepath)\n        for idx,file in enumerate(sub_df['id']):\n            ts = time.time()\n            print('Processingg...',file)\n            filepath = os.path.join(testfolder,file+'.tiff')\n            if not os.path.exists(filepath):\n                continue\n            wsiImage = tiff.imread(filepath)\n            if len(wsiImage.shape) == 5:\n                wsiImage = np.transpose(wsiImage.squeeze(), (1, 2, 0))\n\n            if wsiImage.shape[0]==3:\n                wsiImage = np.transpose(wsiImage, (1, 2, 0))\n\n            wsiMask1 = rle2mask(sub_df['predicted'][idx],\n                                 (wsiImage.shape[1], wsiImage.shape[0]))\n\n            filepath = os.path.join(testfolder, file + '_mask.tiff')\n            print('time taken',time.time()-ts)\n            cv2.imwrite(filepath, wsiMask1)\n            if tryboth:\n                #rle_encode_less_memory\n                h1,w1 = wsiMask1.shape\n                maskhalf = cv2.resize(wsiMask1,(w1,h1))#interpolation=cv2.INTER_AREA)\n                tsb = time.time()\n                code1 = rle_encode_less_memory(maskhalf,img_size_hw=(h1,w1))\n                print('time taken-les memory-', time.time() - tsb)\n\n\n\n                tsb = time.time()\n                code2 = rle_numba_encode_method2(maskhalf, img_size_hw=(h1, w1))\n                print('time taken-2', time.time() - tsb)\n\n                print('if code is same-numba',code1==code2)\n\n\n                if 1:\n                    mask1_re = rle2mask(code1,(w1,h1))\n                    filepath = os.path.join(testfolder, file + '_mask_re1.tiff')\n                    cv2.imwrite(filepath, mask1_re)\n                    mask2_re = rle2mask(code2,(w1,h1))\n                    filepath = os.path.join(testfolder, file + '_mask_re2.tiff')\n                    cv2.imwrite(filepath, mask2_re)\n                    print('are mask same',np.array_equal(mask1_re,mask2_re))\n\n\n\n    def save_train_imgAndMask(self):\n        for idx,file in enumerate(self.train_df['id']):\n            print('Processingg...',file)\n            filepath = os.path.join(self.train_img_folder,file+'.tiff')\n            wsiImage = tiff.imread(filepath)\n\n            if len(wsiImage.shape) == 5:\n                wsiImage = np.transpose(wsiImage.squeeze(), (1, 2, 0))\n\n            if wsiImage.shape[0]==3:\n                wsiImage = np.transpose(wsiImage, (1, 2, 0))\n\n            wsiMask = rle2mask(self.train_df['encoding'][idx],\n                                 (wsiImage.shape[1], wsiImage.shape[0]))\n            wsiMask = wsiMask*255\n\n            print(f'Image Shape: {wsiImage.shape}')\n            print(f'Mask Shape: {wsiMask.shape}')\n\n            if True:\n                filepath = os.path.join(self.save_img_mask, file + f'_mask.tiff')\n                cv2.imwrite(filepath, wsiMask)\n            if self.point_sampling_level !=1:\n                factor = 2**self.point_sampling_level\n                rescaled = (wsiImage.shape[1] // factor, wsiImage.shape[0] // factor)\n                wsiImage = cv2.resize(wsiImage, rescaled,interpolation=cv2.INTER_LINEAR)\n                wsiMask = cv2.resize(wsiMask, rescaled,interpolation=cv2.INTER_AREA)\n\n                filepath = os.path.join(self.save_img_mask,file+f'_resizefactor_{self.point_sampling_level}.tiff')\n                cv2.imwrite(filepath,wsiImage)\n                filepath = os.path.join(self.save_img_mask,file+f'_resizefactor_{self.point_sampling_level}_mask.tiff')\n                cv2.imwrite(filepath, wsiMask)\n                if self.superImposed:\n                    vis_img = wsiImage\n                    vis_img[wsiMask>0, 0] = 0\n                    vis_img[wsiMask>0, 1] = 255\n                    vis_img[wsiMask>0, 2] = 0\n                    filepath = os.path.join(self.save_img_mask,\n                                            file + f'_resizefactor_{self.point_sampling_level}_sumperImposed.tiff')\n                    cv2.imwrite(filepath, vis_img)\n\n    def get_tissue_mask(self,img,method='std',stdthd=5):\n\n        if method=='std':\n            mask = np.std(img,axis=-1) >stdthd\n        else:\n            #otsu\n            if np.ndim(img) == 3:\n                img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)\n            thd = threshold_otsu(img)\n            mask = img < thd\n\n        mask = self.postprocessing(mask)\n        return mask\n\n    def process_patch_locations(self,imgName,idx):\n        self.wsi_list.append(imgName)\n        filepath = os.path.join(self.train_img_folder, imgName + '.tiff')\n        wsiImage = tiff.imread(filepath)\n        if len(wsiImage.shape) == 5:\n            wsiImage = np.transpose(wsiImage.squeeze(), (1, 2, 0))\n        if wsiImage.shape[0] == 3:\n            wsiImage = np.transpose(wsiImage, (1, 2, 0))\n\n        wsiMask = rle2mask(self.train_df['encoding'][idx],\n                           (wsiImage.shape[1], wsiImage.shape[0]))\n        wsiMask = wsiMask * 255\n\n        factor = 2 ** self.point_sampling_level\n        rescaled = (wsiImage.shape[1] // factor, wsiImage.shape[0] // factor)\n        wsiImage_rsz = cv2.resize(wsiImage, rescaled, interpolation=cv2.INTER_LINEAR)\n        #wsiImage_gray = cv2.cvtColor(wsiImage_rsz,cv2.COLOR_RGB2GRAY)\n        valid_region_mask = self.get_tissue_mask(wsiImage_rsz,method='std')\n\n        H,W = valid_region_mask.shape\n\n        if False:\n            mask_filename = os.path.join(self.save_img_mask, f\"{imgName}_tissuemask.png\")\n            tiff.imwrite(mask_filename, 255*np.uint8(valid_region_mask))\n\n        stride_mask_level = self.stride // factor\n        patch_mask_level = self.patch_size // factor\n\n        grid_mask = np.zeros_like(valid_region_mask)\n        x_grid = np.arange(0, W, stride_mask_level)\n        y_grid = np.arange(0, H, stride_mask_level)\n        Xcord, Ycord = np.meshgrid(x_grid, y_grid)\n        grid_mask[Ycord, Xcord] = True\n        yxPair = np.argwhere(grid_mask)\n        area_thd = 0.20*patch_mask_level*patch_mask_level\n        yxpairList=[]\n        for y, x in yxPair:\n            if np.sum(valid_region_mask[y:y + patch_mask_level, x: x + patch_mask_level]) > area_thd:\n                # if np.sum(valid_region_mask[y:y + patch_mask_level, x: x + patch_mask_level]==0) > 8*area_thd:\n                #     continue\n                imgPatch = wsiImage[y*factor:y*factor+self.patch_size,x*factor:x*factor+self.patch_size,:]\n                if np.std(imgPatch) < 10:\n                    continue\n                if len(np.unique(imgPatch)) <25:\n                    continue\n                yxpairList.append((imgName,y * factor, x * factor))\n\n        if self.save_patch:\n            print('Saving patches...')\n            for item in yxpairList:\n                flag = 0\n                imagename,y,x = item\n                img_crop = wsiImage[y:y+self.patch_size,x:x+self.patch_size,:]\n                mask_crop = wsiMask[y:y+self.patch_size,x:x+self.patch_size]\n                mask_crop = cv2.dilate(mask_crop, disk(2), iterations=1)\n                if np.sum(mask_crop)>0:\n                    flag = 1\n                mask_deeplab = np.zeros_like(mask_crop)\n                mask_deeplab[mask_crop==255] = 1\n                mask_dilated = cv2.dilate(mask_crop,self.kernel,iterations=2)\n                mask_boundray = mask_dilated - mask_crop\n                mask_deeplab[mask_boundray>0]=255\n                patch_filename = os.path.join(self.save_train_img,f\"{imagename}_Y_{y}_X_{x}_tn_{flag}.png\")\n                mask_filename = os.path.join(self.save_train_mask,f\"{imagename}_Y_{y}_X_{x}_tn_{flag}.png\")\n                cv2.imwrite(patch_filename,cv2.cvtColor(img_crop,cv2.COLOR_RGB2BGR))\n                cv2.imwrite(mask_filename,mask_deeplab)#cv2.COLOR_RGB2BGR))\n\n                \n        \n        if self.save_gan_pair_list:\n            print('Gan pairpatches...')\n            for item in yxpairList:\n                flag = 0\n                imagename,y,x = item\n                img_crop = wsiImage[y:y+self.patch_size,x:x+self.patch_size,:]\n                mask_crop = wsiMask[y:y+self.patch_size,x:x+self.patch_size]\n                mask_crop = cv2.dilate(mask_crop, disk(2), iterations=1)\n                if np.sum(mask_crop)>0:\n                    flag = 1\n                mask_deeplab = np.zeros_like(mask_crop)\n                mask_deeplab[mask_crop==255] = 1\n                mask_dilated = cv2.dilate(mask_crop,self.kernel,iterations=2)\n                mask_boundray = mask_dilated - mask_crop\n                mask_deeplab[mask_boundray>0]=255\n                patch_filename = os.path.join(self.save_train_img,f\"{imagename}_Y_{y}_X_{x}_tn_{flag}.png\")\n                mask_filename = os.path.join(self.save_train_mask,f\"{imagename}_Y_{y}_X_{x}_tn_{flag}.png\")\n                \n                if np.sum(mask_crop >0) > 0.2*mask_crop.shape[0]*mask_crop.shape[1]:\n                    self.gan_pair_list.append([patch_filename,mask_filename])\n                    cv2.imwrite(patch_filename,cv2.cvtColor(img_crop,cv2.COLOR_RGB2BGR))\n                    cv2.imwrite(mask_filename,mask_deeplab)#cv2.COLOR_RGB2BGR))\n\n        return yxpairList\n\n    def postprocessing(self,img):\n        if img.dtype == np.bool:\n            img = 255* np.uint8(img)\n        kernel = disk(2)\n        # remove unwanted small detection\n        # self.stitch_image = cv2.erode(self.stitch_image,kernel,iterations=1)\n        # remove small white object- opening  ; Closing fills the hole in white image\n        img = cv2.morphologyEx(img, cv2.MORPH_OPEN, kernel, iterations=2)\n        return img\n\n    def save_train_locations(self):\n        yxLocations = []\n        for idx, file in enumerate(self.train_df['id']):\n            print('Processingg...', file)\n            yxLocations += self.process_patch_locations(file,idx)\n\n\n        np.save(os.path.join(self.save_folder,'yxLocations.npy'),yxLocations)\n\n\n\n    def split_train_valid(self,patch_folder=None,folds=None,test_ratio=0.25):\n        if patch_folder is None:\n            patch_folder = self.save_train_img\n        files = list(glob.glob(os.path.join(patch_folder, '*.png')))\n\n        suffix = 'kd_625_'\n        txtfolderpath = os.path.dirname(patch_folder)\n        if folds is None:\n            imageNameList, labelNameList = [], []\n            for file in files:\n                imageNameList.append(file)\n                lablelstr = file.split('.')[0]\n                label = int(lablelstr.split('_')[-1])\n                labelNameList.append(label)\n\n            X_train, X_test, y_train, y_test = train_test_split(imageNameList, labelNameList,\n                                                                test_size=test_ratio, stratify=labelNameList,\n                                                                shuffle=True, random_state=42)\n\n            with open(os.path.join(txtfolderpath, f'train_{suffix}.txt'), 'w') as fp:\n                for file in X_train:\n                    basename = os.path.basename(file)\n                    base = basename.split('.')[0]\n                    fp.write(base)\n                    fp.write('\\n')\n\n            with open(os.path.join(txtfolderpath, f'val_{suffix}.txt'), 'w') as fp:\n                for file in X_test:\n                    basename = os.path.basename(file)\n                    base = basename.split('.')[0]\n                    fp.write(base)\n                    fp.write('\\n')\n\n        else:\n            total_wsi  = len(self.wsi_list)\n            print(self.wsi_list)\n            for key in folds:\n                valid_files = folds[key]\n                X_train ,X_test = [] ,[]\n                for file in files:\n                    basename = os.path.basename(file)\n                    first_name = basename.split('_')[0]\n                    if first_name in valid_files:\n                        X_test.append(file)\n                    else:\n                        X_train.append(file)\n\n\n                with open(os.path.join(txtfolderpath, f'train_{suffix}_{key}.txt'), 'w') as fp:\n                    for file in X_train:\n                        basename = os.path.basename(file)\n                        base = basename.split('.')[0]\n                        fp.write(base)\n                        fp.write('\\n')\n\n                with open(os.path.join(txtfolderpath, f'val_{suffix}_{key}.txt'), 'w') as fp:\n                    for file in X_test:\n                        basename = os.path.basename(file)\n                        base = basename.split('.')[0]\n                        fp.write(base)\n                        fp.write('\\n')\n\n        print('Files are written at-',txtfolderpath)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# !rm -r ./imgdata\n# !rm -r ./mask\n# !rm -r ./patch","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\nimport PIL\nimport sys\nPIL.Image.MAX_IMAGE_PIXELS = 933120000\nts  = time.time()\nROOT_PATH = r'/kaggle/input/hubmap-kidney-segmentation/'\nTRAIN_CSV = os.path.join(ROOT_PATH,'train.csv')\nTRAIN_IMAGE_FOLDER = os.path.join(ROOT_PATH,'train')\nSAVE_FOLDER = r'./'\nGAN_LIST = True\n\nFOLDS_WSI_WISE = {\n         'fold0': ['2f6ecfcdf','aaa6a05cc'],\n         'fold1': ['cb2d976f4', '0486052bb'],\n         'fold2': ['e79de561c', '095bf7a1f','1e2425f28'],\n         }\n\nts = time.time()\nobjData= KDatasetPreparation(train_csv=TRAIN_CSV,\n                             point_sampling_level=3,\n                             save_patch=False,\n                             save_gan_pair_list = GAN_LIST,\n                             patch_size=625,# idear i sto have big an dthen crop 512\n                             save_folder=SAVE_FOLDER,\n                             train_img_folder=TRAIN_IMAGE_FOLDER)\n# objData.get_mask_from_submission(subfilepath=r'C:\\Users\\sjha157731\\Downloads\\test\\submission.csv',\n#                                  testfolder=r'C:\\Users\\sjha157731\\Downloads\\test')\n\n#objData.save_train_imgAndMask()\n#objData.save_train_locations()\n#np.save(os.path.join(SAVE_FOLDER,'gan_pair_list.npy'),objData.gan_pair_list)\n#objData.split_train_valid(test_ratio=0.25,folds=FOLDS_WSI_WISE)\n#print('Total Time for pair genertaion',time.time()-ts)\n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import shutil\npatch_path = r'/kaggle/input/pix2pixmaskimage/img_patches/patch'\nmask_path = r'/kaggle/input/pix2pixmaskimage/mask_patches/mask'\nnpy_path = r'/kaggle/input/pix2pixmaskimage/gan_pair_list.npy'\n#dst = r'./mask'\n#shutil.copytree(mask_path,dst)\n#dst = r'./patch'\n#shutil.copytree(patch_path,dst)\n#shutil.copy(npy_path,'./')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#!zip -r ./img_patches.zip ./patch\n#!zip -r ./mask_patches.zip ./mask","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import numpy as np\n#sanity chekc\npathdata = r'/kaggle/input/pix2pixmaskimage/gan_pair_list.npy'\ndata = np.load(pathdata,allow_pickle=True)\ndata = np.array(data)\nprint(data.shape)\n\n#updated path\nupdated = []\nfor i in range(data.shape[0]):\n    item = data[i]\n    x_c = item[0].replace('./','/kaggle/input/pix2pixmaskimage/img_patches/')\n    y_c = item[1].replace('./','/kaggle/input/pix2pixmaskimage/mask_patches/')\n    updated.append([x_c,y_c])\nupdated = np.array(updated)\nnp.save('./gan_pair_list_updated.npy',updated)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**PIX2PIX GAN - Define model **"},{"metadata":{"trusted":true},"cell_type":"code","source":"import numpy as np\nimport tensorflow as tf\nimport os\nimport matplotlib.pyplot as plt\nimport time\n'''\n#https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/tutorials/generative/pix2pix.ipynb#scrollTo=dIbRPFzjmV85\n#https://github.com/tensorflow/docs/blob/master/site/en/tutorials/generative/pix2pix.ipynb\n#https://machinelearningmastery.com/what-are-generative-adversarial-networks-gans/\n#define models\n#Genereator - a Unet - Discriminator - patchGan\n#Kernel size = 4, stride =2\nTaken from paper-\n    Generator - UNET-\n        *Encoder:\n            C64-C128-C256-C512-C512-C512-C512-C512\n        *Decoder:\n            CD512-CD512-CD512-C512-C256-C128-C64\n        *   (D/E)Let Ck denote a Convolution-BatchNorm-ReLU layer with k filters. CDk denotes a \n            Convolution-BatchNorm-Dropout-ReLU layer with a dropout rate of 50%. \n        *   (D/E)All convolutions are 4x4 spatial filters applied with stride 2. \n            Convolutions in the encoder, and in the discriminator, downsample\n            by a factor of 2, whereas in the decoder they upsample by a\n            factor of 2.\n        *   (E)All ReLUs in the encoder are leaky, with slope 0.2, \n        *   (E)ReLUs in the decoder are not leaky.\n        *   (D/E)Weights were initialized from a Gaussian distribution with mean 0 and standard deviation 0.02\n        *   (E)Batch-Norm is not applied to the first C64 layer in the encoder\n        *   (D)After the last layer in the decoder, a convolution is applied to map to the number of output channels (3 in general,\n            except in colorization, where it is 2), followed by a Tanh function\n    \n    Discriminator- patchGAN\n         *  The 70 x70 discriminator architecture is:  C64-C128-C256-C512\n            After the last layer, a convolution is applied to map to\n            a 1-dimensional output, followed by a Sigmoid function.\n        *   As an exception to the above notation, BatchNorm is not\n            applied to the first C64 layer. All ReLUs are leaky, with\n            slope 0.2\n'''\ndef DownsampleBlock(nFilter,kernel=4,stride=2,apply_batchnorm=True,apply_dropout=False):\n    #Conv->BN->Leaky-Relu\n    initializer = tf.random_normal_initializer(0., 0.02)\n    block = tf.keras.Sequential()\n    block.add(tf.keras.layers.Conv2D(filters=nFilter,kernel_size=kernel,strides=stride,\n                                     padding='same',kernel_initializer=initializer,\n                                     use_bias=False))\n\n    if apply_batchnorm:\n        block.add(tf.keras.layers.BatchNormalization())\n\n    if apply_dropout:\n        block.add(tf.keras.layers.Dropout(alpha=0.5))\n\n    block.add(tf.keras.layers.LeakyReLU(alpha=0.2))\n    return block\n\ndef UpsampleBlock(nFilter, kernel=4, stride=2, apply_batchnorm=True, apply_dropout=False):\n    # Conv->BN->Dropout->Relu\n    initializer = tf.random_normal_initializer(0., 0.02)\n    block = tf.keras.Sequential()\n    block.add(tf.keras.layers.Conv2DTranspose(filters=nFilter, kernel_size=kernel, strides=stride,\n                                     padding='same', kernel_initializer=initializer,\n                                     use_bias=False))\n\n    if apply_batchnorm:\n        block.add(tf.keras.layers.BatchNormalization())\n\n    if apply_dropout:\n        block.add(tf.keras.layers.Dropout(rate=0.5))\n\n    block.add(tf.keras.layers.ReLU())\n\n    return block\n\n\ndef generator_model(input_shape=[256,256,3]):\n    nFilters = 64\n    nUpFilters = 512 #(64*8)\n    OUTPUT_CHANNELS = input_shape[2]\n    input = tf.keras.layers.Input(shape=input_shape)\n\n    #encoder -C64-C128-C256-C512-C512-C512-C512-C512\n    c064 = DownsampleBlock(nFilter=nFilters,apply_batchnorm=False,apply_dropout=False)(input) #-1,128x128x64\n    c128 = DownsampleBlock(nFilter=nFilters*2,apply_dropout=False)(c064)   #-1,64x64x128\n    c256 = DownsampleBlock(nFilter=nFilters*4,apply_dropout=False)(c128)   #-1,32x32x256\n    c0x512 = DownsampleBlock(nFilter=nFilters*8,apply_dropout=False)(c256)   #-1,16x16x512\n    c1x512 = DownsampleBlock(nFilter=nFilters*8,apply_dropout=False)(c0x512)   #-1,8x8x512\n    c2x512 = DownsampleBlock(nFilter=nFilters*8,apply_dropout=False)(c1x512)   #-1,4x4x512\n    c3x512 = DownsampleBlock(nFilter=nFilters*8,apply_dropout=False)(c2x512)   #-1,2x2x512\n    c4x512 = DownsampleBlock(nFilter=nFilters*8,apply_dropout=False)(c3x512)   #-1,1x1x512\n\n    #decoder - CD512-CD512-CD512-C512-C256-C128-C64\n    u0512 = UpsampleBlock(nFilter=nUpFilters,apply_dropout=True)(c4x512) #-1,2x2x512\n    cat_1 = tf.keras.layers.Concatenate()([u0512, c3x512])               #-1,2x2x2048\n    u1512 = UpsampleBlock(nFilter=nUpFilters,apply_dropout=True)(cat_1)  #-1,4x4x1024\n    cat_2 = tf.keras.layers.Concatenate()([u1512, c2x512])               #-1,4x4x2048\n    u2512 = UpsampleBlock(nFilter=nUpFilters, apply_dropout=True)(cat_2) # -1,8x8x1024\n    cat_3 = tf.keras.layers.Concatenate()([u2512, c1x512])               # -1,8x8x2048\n    u3512 = UpsampleBlock(nFilter=nUpFilters, apply_dropout=False)(cat_3)  # -1,16x16x1024\n    cat_4 = tf.keras.layers.Concatenate()([u3512, c0x512])                 # -1,16x16x2048\n\n    u4512 = UpsampleBlock(nFilter=nUpFilters//2, apply_dropout=False)(cat_4)   # -1,32x32x256\n    cat_5 = tf.keras.layers.Concatenate()([u4512, c256])                       # -1,32x32x512\n\n    u5512 = UpsampleBlock(nFilter=nUpFilters // 4, apply_dropout=False)(cat_5)  # -1,64x64x128\n    cat_6 = tf.keras.layers.Concatenate()([u5512, c128])                        # -1,64x64x256\n\n    u6512 = UpsampleBlock(nFilter=nUpFilters // 8, apply_dropout=False)(cat_6)  # -1,128x128x64\n    cat_7 = tf.keras.layers.Concatenate()([u6512, c064])                        #-1,128x128x128\n\n    #Exception for last layer\n    initializer = tf.random_normal_initializer(0., 0.02)\n    last = tf.keras.layers.Conv2DTranspose(OUTPUT_CHANNELS, 4,\n                                           strides=2,\n                                           padding='same',\n                                           kernel_initializer=initializer,\n                                           activation='tanh') (cat_7)           # -1x256x256x3\n\n    model = tf.keras.Model(inputs=input, outputs=last)\n    return model\n\ndef discriminator_model(input_shape=[256,256,3],outlogit_no_sigmoid=True):\n    '''\n    This model takes 2 inouts - src & target\n    Each 30x30 patch of the output classifies a 70x70 portion of the\n    input image (such an architecture is called a PatchGAN).\n    '''\n    initializer = tf.random_normal_initializer(0., 0.02)\n    inp = tf.keras.layers.Input(shape=input_shape, name='input_image')\n    tar = tf.keras.layers.Input(shape=input_shape, name='target_image')\n\n    catx = tf.keras.layers.concatenate([inp, tar])  # (-1x256x256xchannels*2)\n    x = DownsampleBlock(nFilter=64,apply_batchnorm=False,apply_dropout=False)(catx) #-1x128x128x64\n    x = DownsampleBlock(nFilter=128,apply_dropout=False)(x)     #-1x64x64x128\n    x = DownsampleBlock(nFilter=256,apply_dropout=False)(x)     #-1x32x32x256\n\n    #taken from Tensorflow Colab\n    pad = tf.keras.layers.ZeroPadding2D()(x)            #-1x34x34x256\n    conv = tf.keras.layers.Conv2D(512, 4, strides=1,\n                                  kernel_initializer=initializer,\n                                  use_bias=False)(pad)  #-1x31x31x512\n\n    bn = tf.keras.layers.BatchNormalization()(conv)\n\n    leaky_relu = tf.keras.layers.LeakyReLU(alpha=0.2)(bn)\n\n    zero_pad2 = tf.keras.layers.ZeroPadding2D()(leaky_relu)  #-1x33x33x512\n\n    #as per paper\n    last = tf.keras.layers.Conv2D(1, 4, strides=1,\n                                  kernel_initializer=initializer)(zero_pad2)  #-1x30x30x1\n    if not outlogit_no_sigmoid:\n        last = tf.keras.layers.Activation(activation='sigmoid')(last)\n\n    return tf.keras.Model(inputs=[inp, tar], outputs=last)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def load_image_train(image_file):\n    input_image, real_image = load_local(image_file)\n    input_image, real_image = random_jitter(input_image, real_image)\n    input_image, real_image = normalize(input_image, real_image)\n    return input_image, real_image\n\ndef load_image_test(image_file):\n    input_image, real_image = load_local(image_file)\n    input_image, real_image = resize(input_image, real_image,\n                                   IMG_HEIGHT, IMG_WIDTH)\n    input_image, real_image = normalize(input_image, real_image)\n    return input_image, real_image\n\ndef load(image_file):\n    image = tf.io.read_file(image_file)\n    image = tf.image.decode_jpeg(image)\n\n    w = tf.shape(image)[1]\n\n    w = w // 2\n    real_image = image[:, :w, :]\n    input_image = image[:, w:, :]\n\n    input_image = tf.cast(input_image, tf.float32)\n    real_image = tf.cast(real_image, tf.float32)\n\n    return input_image, real_image\n\ndef load_local(image_file):\n    #0: Orig image /1:- Mask\n    image = tf.io.read_file(image_file[0])\n    image = tf.image.decode_png(image)\n    real_image = tf.cast(image, tf.float32)\n    image = tf.io.read_file(image_file[1])\n    image = tf.image.decode_png(image)\n    #image = tf.print(image,[image.shape])\n    if True:\n          op1 = tf.zeros_like(image)  # 0\n          op2 = op1 + 255  # 255\n          image = tf.where(tf.equal(image, tf.constant(255, dtype=tf.uint8)), op1, image)\n          #image = tf.where(tf.equal(image, tf.constant(1, dtype=tf.uint8)), op2, label)\n    image = tf.image.grayscale_to_rgb(image)\n    input_image = tf.cast(image, tf.float32)\n\n    return input_image, real_image\n\ndef resize(input_image, real_image, height, width):\n    input_image = tf.image.resize(input_image, [height, width],\n                                method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)\n    real_image = tf.image.resize(real_image, [height, width],\n                               method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)\n\n    return input_image, real_image\n\n\ndef random_crop(input_image, real_image):\n    stacked_image = tf.stack([input_image, real_image], axis=0)\n    cropped_image = tf.image.random_crop(\n      stacked_image, size=[2, IMG_HEIGHT, IMG_WIDTH, 3])\n\n    return cropped_image[0], cropped_image[1]\n# normalizing the images to [-1, 1]\n\ndef normalize(input_image, real_image):\n    #input_image = (input_image / 127.5) - 1 # already -0/1 range\n    real_image = (real_image / 127.5) - 1\n\n    return input_image, real_image\n\n@tf.function()\ndef random_jitter(input_image, real_image):\n    # resizing to 286 x 286 x 3\n    input_image, real_image = resize(input_image, real_image, 286, 286)\n\n    # randomly cropping to 256 x 256 x 3\n    input_image, real_image = random_crop(input_image, real_image)\n\n    if tf.random.uniform(()) > 0.5:\n        # random mirroring\n        input_image = tf.image.flip_left_right(input_image)\n        real_image = tf.image.flip_left_right(real_image)\n\n    return input_image, real_image","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def generator_loss(discriminator_out_gen_in,generator_out,target):\n    '''\n    Idea is that discriminator shld predict real for photo generated\n    from generator for real training data\n    :param discriminator_out_gen_in:\n    :param generator_out:\n    :param target:\n    '''\n    LAMBDA = 100\n    #to maintain structure\n    gen_L1_loss = tf.reduce_mean(tf.abs(target - generator_out))\n    #bce loss-\n    gt_label = tf.ones_like(discriminator_out_gen_in)\n    gen_bce_loss = tf.keras.losses.BinaryCrossentropy(from_logits=True)(gt_label,discriminator_out_gen_in) #y_true, y_pred\n    gen_total_loss = gen_bce_loss + LAMBDA*gen_L1_loss\n\n    return gen_total_loss,gen_bce_loss,gen_L1_loss\n\ndef discriminator_loss(disc_real_output, disc_generated_output):\n    loss_object = tf.keras.losses.BinaryCrossentropy(from_logits=True)\n    real_loss = loss_object(tf.ones_like(disc_real_output), disc_real_output)\n    generated_loss = loss_object(tf.zeros_like(disc_generated_output), disc_generated_output)\n    total_disc_loss = real_loss + generated_loss\n    #total_disc_loss = 0.5*(real_loss + generated_loss) # from paper\n\n    return total_disc_loss","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"discriminator = discriminator_model(outlogit_no_sigmoid=True)\ntf.keras.utils.plot_model(discriminator, to_file='./dis_model.png',show_shapes=True, dpi=64)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"generator = generator_model()\ntf.keras.utils.plot_model(generator, to_file='./gen_model.png',show_shapes=True, dpi=64)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def generate_images(model, test_input, tar):\n  prediction = model(test_input, training=True)\n  plt.figure(figsize=(15,15))\n\n  display_list = [test_input[0], tar[0], prediction[0]]\n  title = ['Input Image', 'Ground Truth', 'Predicted Image']\n\n  for i in range(3):\n    plt.subplot(1, 3, i+1)\n    plt.title(title[i])\n    # getting the pixel values between [0, 1] to plot it.\n    plt.imshow(display_list[i] * 0.5 + 0.5)\n    plt.axis('off')\n  plt.show()\n\n@tf.function\ndef train_step(input_image, target, epoch):\n  with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:\n    #get generator out -\n    gen_output = generator(input_image, training=True)\n\n    disc_real_output = discriminator([input_image, target], training=True)\n    disc_generated_output = discriminator([input_image, gen_output], training=True)\n\n    gen_total_loss, gen_gan_loss, gen_l1_loss = generator_loss(disc_generated_output, gen_output, target)\n    disc_loss = discriminator_loss(disc_real_output, disc_generated_output)\n\n  generator_gradients = gen_tape.gradient(gen_total_loss,\n                                          generator.trainable_variables)\n  discriminator_gradients = disc_tape.gradient(disc_loss,\n                                               discriminator.trainable_variables)\n\n  generator_optimizer.apply_gradients(zip(generator_gradients,\n                                          generator.trainable_variables))\n  discriminator_optimizer.apply_gradients(zip(discriminator_gradients,\n                                              discriminator.trainable_variables))\n\n  with summary_writer.as_default():\n    tf.summary.scalar('gen_total_loss', gen_total_loss, step=epoch)\n    tf.summary.scalar('gen_gan_loss', gen_gan_loss, step=epoch)\n    tf.summary.scalar('gen_l1_loss', gen_l1_loss, step=epoch)\n    tf.summary.scalar('disc_loss', disc_loss, step=epoch)\n\n\ndef fit(train_ds, epochs, test_ds):\n  for epoch in range(epochs):\n    start = time.time()\n    display.clear_output(wait=True)\n    #plt.clf()\n\n    for example_input, example_target in test_ds.take(1):\n      generate_images(generator, example_input, example_target)#epoch=epoch)\n    print(\"Epoch: \", epoch)\n\n    # Train\n    for n, (input_image, target) in train_ds.enumerate():\n      print('.', end='')\n      if (n+1) % 100 == 0:\n        print()\n      train_step(input_image, target, epoch)\n    print()\n\n    # saving (checkpoint) the model every 10 epochs\n    if (epoch + 1) % 5 == 0:\n      checkpoint.save(file_prefix = checkpoint_prefix)\n\n    print ('Time taken for epoch {} is {} sec\\n'.format(epoch + 1,\n                                                        time.time()-start))\n  checkpoint.save(file_prefix = checkpoint_prefix)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Load Training Data**"},{"metadata":{"trusted":true},"cell_type":"code","source":"from matplotlib import pyplot as plt\nfrom IPython import display\n%matplotlib inline","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"BUFFER_SIZE = 400\nBATCH_SIZE = 16\nEPOCHS = 75\nIMG_WIDTH = 256\nIMG_HEIGHT = 256\nPLOT = False\nPLOT_JITTER = False\ndo_test = False\nimport sys\nfrom sklearn.model_selection import train_test_split\npathdata = r'./gan_pair_list_updated.npy'\ndata = np.load(pathdata,allow_pickle=True)\ndata = np.array(data)\ntrain_files, test_files, = train_test_split(data,test_size=0.20,\n                                    shuffle=True,random_state=42)\nBUFFER_SIZE_TRAIN = train_files.shape[0]\nBUFFER_SIZE_TEST = test_files.shape[0]\nprint(data.shape,train_files.shape,test_files.shape)\n\nif True:\n    train_dataset = tf.data.Dataset.from_tensor_slices(train_files)\n    train_dataset = train_dataset.map(load_image_train,\n                                      num_parallel_calls=tf.data.experimental.AUTOTUNE)\n    train_dataset = train_dataset.shuffle(400)\n    train_dataset = train_dataset.batch(1)\n    #(img1,img2) = train_dataset.take(1)\n    \n    for input_image, target in train_dataset.as_numpy_iterator():\n        plt.subplot(121)\n        plt.imshow(np.squeeze(input_image))\n        plt.subplot(122)\n        plt.imshow(np.squeeze(target))\n        plt.show()\n        break\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Define generators**"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dataset = tf.data.Dataset.from_tensor_slices(train_files)\n#train_dataset = tf.data.Dataset.list_files(os.path.join(PATH , 'train/*.jpg'))\ntrain_dataset = train_dataset.map(load_image_train,\n                                  num_parallel_calls=tf.data.experimental.AUTOTUNE)\ntrain_dataset = train_dataset.shuffle(BUFFER_SIZE_TRAIN)\ntrain_dataset = train_dataset.batch(BATCH_SIZE)\n\ntest_dataset = tf.data.Dataset.from_tensor_slices(test_files)\n#test_dataset = tf.data.Dataset.list_files(os.path.join(PATH , 'test/*.jpg'))\ntest_dataset = test_dataset.map(load_image_test)\ntest_dataset = test_dataset.batch(BATCH_SIZE)\n\n\n\ngenerator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\ndiscriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n\ncheckpoint_dir = './kd_training_checkpoints'\ncheckpoint_prefix = os.path.join(checkpoint_dir, \"ckpt\")\ncheckpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer,\n                                 discriminator_optimizer=discriminator_optimizer,\n                                 generator=generator,\n                                 discriminator=discriminator)\n\n\nlog_dir = \"logs/\"\nsummary_writer = tf.summary.create_file_writer(\n    log_dir + \"fit/\" + datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\"))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%load_ext tensorboard\n%tensorboard --logdir {log_dir}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Training Loop**"},{"metadata":{"trusted":true},"cell_type":"code","source":"\nfor example_input, example_target in test_dataset.take(1):\n    generate_images(generator, example_input, example_target)\n\nif True:\n    try:\n        fit(train_dataset, EPOCHS, test_dataset)\n    except Exception as E:\n        print(E)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if True:\n    checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))\n    # Run the trained model on a few examples from the test dataset\n    cont =0\n    for inp, tar in test_dataset.take(25):\n        generate_images(generator, inp, tar)\n        cont +=1","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"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":4,"nbformat_minor":4}