{"cells":[{"metadata":{},"cell_type":"markdown","source":"# FAST STYLE TRANSFER"},{"metadata":{},"cell_type":"markdown","source":"<div>\n  <img src='https://github.com/tarun-bisht/fast-style-transfer/raw/master/data/images/style.jpg' height=\"346px\">\n  <img src='https://github.com/tarun-bisht/fast-style-transfer/raw/master/data/images/content.jpg' height=\"346px\">\n  <img src='https://github.com/tarun-bisht/fast-style-transfer/raw/master/output/styled.jpg' height=\"512px\">\n</div>"},{"metadata":{},"cell_type":"markdown","source":"Stylize any photo or video in style of famous paintings using Neural Style Transfer.\n- This is hundreds of times faster than the optimization-based method presented by [Gatys et al](https://arxiv.org/abs/1508.06576) so called fast style transfer.\n- We train a feedforward network that apply artistic styles to images using loss function defined in [Gatys et al](https://arxiv.org/abs/1508.06576) paper.\n- Feed forward network is a residual autoencoder network that takes content image as input and spits out stylized image.\n- Model also uses instance normalization instead of batch normalization based on the paper [Instance Normalization: The Missing Ingredient for Fast Stylization](https://arxiv.org/abs/1607.08022)\n- Training is done by using perceptual loss defined in paper [Perceptual Losses for Real-Time Style Transfer and Super-Resolution](https://arxiv.org/abs/1603.08155).\n- Vgg19 is used to calculate perceptual loss more working described on paper."},{"metadata":{},"cell_type":"markdown","source":"## SETTING UP"},{"metadata":{},"cell_type":"markdown","source":"### Importing all dependencies"},{"metadata":{"id":"wTSIrk91pMXl","trusted":true},"cell_type":"code","source":"import numpy as np\nimport tensorflow as tf\nfrom tensorflow.keras.applications import vgg19\nfrom tensorflow.keras.models import load_model,Model\nfrom PIL import Image\nimport time\nimport matplotlib.pyplot as plt\nimport matplotlib\nimport requests\nimport base64\nimport os\nfrom pathlib import Path\nfrom io import BytesIO\nmatplotlib.rcParams['figure.figsize'] = (12,12)\nmatplotlib.rcParams['axes.grid'] = False","execution_count":null,"outputs":[]},{"metadata":{"id":"uxV3zaDLc7wo"},"cell_type":"markdown","source":"### Utility Functions"},{"metadata":{"id":"UCfbjcUnpTzP","trusted":true},"cell_type":"code","source":"def load_image(image_path,dim=None,resize=False):\n    img= Image.open(image_path)\n    if dim:\n        if resize:\n            img=img.resize(dim)\n        else:\n            img.thumbnail(dim)\n    img= img.convert(\"RGB\")\n    return np.array(img)","execution_count":null,"outputs":[]},{"metadata":{"id":"PfmBkMqk7S6v","trusted":true},"cell_type":"code","source":"def load_url_image(url,dim=None,resize=False):\n    img_request=requests.get(url)\n    img= Image.open(BytesIO(img_request.content))\n    if dim:\n        if resize:\n            img=img.resize(dim)\n        else:\n            img.thumbnail(dim)\n    img= img.convert(\"RGB\")\n    return np.array(img)","execution_count":null,"outputs":[]},{"metadata":{"id":"jB4-cSv6pbCI","trusted":true},"cell_type":"code","source":"def array_to_img(array):\n    array=np.array(array,dtype=np.uint8)\n    if np.ndim(array)>3:\n        assert array.shape[0]==1\n        array=array[0]\n    return Image.fromarray(array)","execution_count":null,"outputs":[]},{"metadata":{"id":"O-wiqru1pcvn","trusted":true},"cell_type":"code","source":"def show_image(image,title=None):\n    if len(image.shape)>3:\n        image=tf.squeeze(image,axis=0)\n    plt.imshow(image)\n    if title:\n        plt.title=title","execution_count":null,"outputs":[]},{"metadata":{"id":"clS-9PcJMhbi","trusted":true},"cell_type":"code","source":"def plot_images_grid(images,num_rows=1):\n    n=len(images)\n    if n > 1:\n        num_cols=np.ceil(n/num_rows)\n        fig,axes=plt.subplots(ncols=int(num_cols),nrows=int(num_rows))\n        axes=axes.flatten()\n        fig.set_size_inches((15,15))\n        for i,image in enumerate(images):\n            axes[i].imshow(image)\n    else:\n        plt.figure(figsize=(10,10))\n        plt.imshow(images[0])","execution_count":null,"outputs":[]},{"metadata":{"id":"SW3vJl1DJhpZ"},"cell_type":"markdown","source":"## SETTING UP LOSS MODEL\n\nHere pretrained VGG19 model was used to calculate perceptual loss as describe in this [paper](https://arxiv.org/abs/1603.08155). \n\nIf you know about Gatys style transfer then it is same, we need to calculate style and content loss using a pretrained model using them we create total loss. \n\nThis loss is calculated by passing style and content images both into this VGG network and using some layers of network to extract features. Higher layer learns complex features which also preserves some content of image (like higher layers can represent face of dog but lower layers represent lower features like eyes, nose, ears etc.), so for content loss higher layers of networks are used. For style loss lower layers of networks are used since they have learned small features like edges, countours etc. that can imitate brush stroke from a painting. We use multiple layers for style because every layer learns differnt strokes and to look realistic we need some variations in stroke. "},{"metadata":{"id":"eTYybI5UpgzZ","outputId":"406adbca-b463-43cd-83ed-9eb52a122768","trusted":true},"cell_type":"code","source":"vgg=vgg19.VGG19(weights='imagenet',include_top=False)\nvgg.summary()","execution_count":null,"outputs":[]},{"metadata":{"id":"MLmnPM96JnTG"},"cell_type":"markdown","source":"### Defining content and style layers from pretrained model's layers"},{"metadata":{"id":"NWON8T6oplk4","trusted":true},"cell_type":"code","source":"content_layers=['block4_conv2']\nstyle_layers=['block1_conv1',\n            'block2_conv1',\n            'block3_conv1',\n            'block4_conv1',\n            'block5_conv1']\ncontent_layers_weights=[1]\nstyle_layers_weights=[1]*5","execution_count":null,"outputs":[]},{"metadata":{"id":"e2zdmMkTJwAQ"},"cell_type":"markdown","source":"### Creating Loss model which is used to calcuate perceptual loss by combining style and content loss "},{"metadata":{"id":"BV7yF8XefbfF","trusted":true},"cell_type":"code","source":"class LossModel:\n    def __init__(self,pretrained_model,content_layers,style_layers):\n        self.model=pretrained_model\n        self.content_layers=content_layers\n        self.style_layers=style_layers\n        self.loss_model=self.get_model()\n\n    def get_model(self):\n        self.model.trainable=False\n        layer_names=self.style_layers + self.content_layers\n        outputs=[self.model.get_layer(name).output for name in layer_names]\n        new_model=Model(inputs=self.model.input,outputs=outputs)\n        return new_model\n    \n    def get_activations(self,inputs):\n        inputs=inputs*255.0\n        style_length=len(self.style_layers)\n        outputs=self.loss_model(vgg19.preprocess_input(inputs))\n        style_output,content_output=outputs[:style_length],outputs[style_length:]\n        content_dict={name:value for name,value in zip(self.content_layers,content_output)}\n        style_dict={name:value for name,value in zip(self.style_layers,style_output)}\n        return {'content':content_dict,'style':style_dict}","execution_count":null,"outputs":[]},{"metadata":{"id":"UoCALQjMpoDN","trusted":true},"cell_type":"code","source":"loss_model=LossModel(vgg,content_layers,style_layers)","execution_count":null,"outputs":[]},{"metadata":{"id":"h-wwiZZYJ3k3"},"cell_type":"markdown","source":"### Defining loss functions"},{"metadata":{"id":"B7xre5cqp0cd","trusted":true},"cell_type":"code","source":"def content_loss(placeholder,content,weight):\n    assert placeholder.shape == content.shape\n    return weight*tf.reduce_mean(tf.square(placeholder-content))","execution_count":null,"outputs":[]},{"metadata":{"id":"c-OGwBbbp2zM","trusted":true},"cell_type":"code","source":"def gram_matrix(x):\n    gram=tf.linalg.einsum('bijc,bijd->bcd', x, x)\n    return gram/tf.cast(x.shape[1]*x.shape[2]*x.shape[3],tf.float32)","execution_count":null,"outputs":[]},{"metadata":{"id":"Gibfnt2Xp4eS","trusted":true},"cell_type":"code","source":"def style_loss(placeholder,style,weight):\n    assert placeholder.shape == style.shape\n    s=gram_matrix(style)\n    p=gram_matrix(placeholder)\n    return weight*tf.reduce_mean(tf.square(s-p))","execution_count":null,"outputs":[]},{"metadata":{"id":"6DA4K8N7p6hL","trusted":true},"cell_type":"code","source":"def preceptual_loss(predicted_activations,content_activations,style_activations,content_weight,style_weight,content_layers_weights,style_layer_weights):\n    pred_content=predicted_activations[\"content\"]\n    pred_style=predicted_activations[\"style\"]\n    c_loss=tf.add_n([content_loss(pred_content[name],content_activations[name],content_layers_weights[i]) for i,name in enumerate(pred_content.keys())])\n    c_loss=c_loss*content_weight\n    s_loss=tf.add_n([style_loss(pred_style[name],style_activations[name],style_layer_weights[i]) for i,name in enumerate(pred_style.keys())])\n    s_loss=s_loss*style_weight\n    return c_loss+s_loss","execution_count":null,"outputs":[]},{"metadata":{"id":"9JD5R3dNJ7fb"},"cell_type":"markdown","source":"# FAST STYLE TRANSFER MODEL\n\nIt is a encoder-decoder architecture with residual layers. Input images are passed to encoder part and it propogates to decoder part of same size as input and predict generated image. For training this generated image is passed to our loss model (VGG19) and features from different layers were extracted (content layers and style layers) these features are then used to calculate style loss and content loss, whose weighted sum produce perceptual loss that trains the network. The below image from paper describe it well.\n\n![https://arxiv.org/abs/1603.08155](https://miro.medium.com/max/1574/1*Um82GJ99gauIOh0U-S11hQ.png)\n\nThe main highlights of network:\n\n- Residual Layers\n- Encoder Decoder Model\n- output from decoder is passed to loss model(VGG) to alculate loss and train"},{"metadata":{},"cell_type":"markdown","source":"### Defining model layers"},{"metadata":{"id":"sShUKZjAjdh0","trusted":true},"cell_type":"code","source":"class ReflectionPadding2D(tf.keras.layers.Layer):\n    def __init__(self, padding=(1, 1), **kwargs):\n        super(ReflectionPadding2D, self).__init__(**kwargs)\n        self.padding = tuple(padding)\n    def call(self, input_tensor):\n        padding_width, padding_height = self.padding\n        return tf.pad(input_tensor, [[0,0], [padding_height, padding_height], [padding_width, padding_width], [0,0] ], 'REFLECT')","execution_count":null,"outputs":[]},{"metadata":{"id":"tjVTLcrsY_hJ","trusted":true},"cell_type":"code","source":"class InstanceNormalization(tf.keras.layers.Layer):\n    def __init__(self,**kwargs):\n        super(InstanceNormalization, self).__init__(**kwargs)\n    def call(self,inputs):\n        batch, rows, cols, channels = [i for i in inputs.get_shape()]\n        mu, var = tf.nn.moments(inputs, [1,2], keepdims=True)\n        shift = tf.Variable(tf.zeros([channels]))\n        scale = tf.Variable(tf.ones([channels]))\n        epsilon = 1e-3\n        normalized = (inputs-mu)/tf.sqrt(var + epsilon)\n        return scale * normalized + shift","execution_count":null,"outputs":[]},{"metadata":{"id":"uQXmV9GOLuS_","trusted":true},"cell_type":"code","source":"class ConvLayer(tf.keras.layers.Layer):\n    def __init__(self,filters,kernel_size,strides=1,**kwargs):\n        super(ConvLayer,self).__init__(**kwargs)\n        self.padding=ReflectionPadding2D([k//2 for k in kernel_size])\n        self.conv2d=tf.keras.layers.Conv2D(filters,kernel_size,strides)\n        self.bn=InstanceNormalization()\n    def call(self,inputs):\n        x=self.padding(inputs)\n        x=self.conv2d(x)\n        x=self.bn(x)\n        return x","execution_count":null,"outputs":[]},{"metadata":{"id":"-8FfGGGW0drX","trusted":true},"cell_type":"code","source":"class ResidualLayer(tf.keras.layers.Layer):\n    def __init__(self,filters,kernel_size,**kwargs):\n        super(ResidualLayer,self).__init__(**kwargs)\n        self.conv2d_1=ConvLayer(filters,kernel_size)\n        self.conv2d_2=ConvLayer(filters,kernel_size)\n        self.relu=tf.keras.layers.ReLU()\n        self.add=tf.keras.layers.Add()\n    def call(self,inputs):\n        residual=inputs\n        x=self.conv2d_1(inputs)\n        x=self.relu(x)\n        x=self.conv2d_2(x)\n        x=self.add([x,residual])\n        return x","execution_count":null,"outputs":[]},{"metadata":{"id":"NF0TyLpqHfaP","trusted":true},"cell_type":"code","source":"class UpsampleLayer(tf.keras.layers.Layer):\n    def __init__(self,filters,kernel_size,strides=1,upsample=2,**kwargs):\n        super(UpsampleLayer,self).__init__(**kwargs)\n        self.upsample=tf.keras.layers.UpSampling2D(size=upsample)\n        self.padding=ReflectionPadding2D([k//2 for k in kernel_size])\n        self.conv2d=tf.keras.layers.Conv2D(filters,kernel_size,strides)\n        self.bn=InstanceNormalization()\n    def call(self,inputs):\n        x=self.upsample(inputs)\n        x=self.padding(x)\n        x=self.conv2d(x)\n        return self.bn(x)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Constructing style transfer model"},{"metadata":{"id":"_4HpvuO-p9Ao","trusted":true},"cell_type":"code","source":"class StyleTransferModel(tf.keras.Model):\n    def __init__(self,**kwargs):\n        super(StyleTransferModel, self).__init__(name='StyleTransferModel',**kwargs)\n        self.conv2d_1= ConvLayer(filters=32,kernel_size=(9,9),strides=1,name=\"conv2d_1_32\")\n        self.conv2d_2= ConvLayer(filters=64,kernel_size=(3,3),strides=2,name=\"conv2d_2_64\")\n        self.conv2d_3= ConvLayer(filters=128,kernel_size=(3,3),strides=2,name=\"conv2d_3_128\")\n        self.res_1=ResidualLayer(filters=128,kernel_size=(3,3),name=\"res_1_128\")\n        self.res_2=ResidualLayer(filters=128,kernel_size=(3,3),name=\"res_2_128\")\n        self.res_3=ResidualLayer(filters=128,kernel_size=(3,3),name=\"res_3_128\")\n        self.res_4=ResidualLayer(filters=128,kernel_size=(3,3),name=\"res_4_128\")\n        self.res_5=ResidualLayer(filters=128,kernel_size=(3,3),name=\"res_5_128\")\n        self.deconv2d_1= UpsampleLayer(filters=64,kernel_size=(3,3),name=\"deconv2d_1_64\")\n        self.deconv2d_2= UpsampleLayer(filters=32,kernel_size=(3,3),name=\"deconv2d_2_32\")\n        self.deconv2d_3= ConvLayer(filters=3,kernel_size=(9,9),strides=1,name=\"deconv2d_3_3\")\n        self.relu=tf.keras.layers.ReLU()\n    def call(self, inputs):\n        x=self.conv2d_1(inputs)\n        x=self.relu(x)\n        x=self.conv2d_2(x)\n        x=self.relu(x)\n        x=self.conv2d_3(x)\n        x=self.relu(x)\n        x=self.res_1(x)\n        x=self.res_2(x)\n        x=self.res_3(x)\n        x=self.res_4(x)\n        x=self.res_5(x)\n        x=self.deconv2d_1(x)\n        x=self.relu(x)\n        x=self.deconv2d_2(x)\n        x=self.relu(x)\n        x=self.deconv2d_3(x)\n        x = (tf.nn.tanh(x) + 1) * (255.0 / 2)\n        return x\n    \n    ## used to print shapes of each layer to check if input shape == output shape\n    ## I don't know any better solution to this right now\n    def print_shape(self,inputs):\n        print(inputs.shape)\n        x=self.conv2d_1(inputs)\n        print(x.shape)\n        x=self.relu(x)\n        x=self.conv2d_2(x)\n        print(x.shape)\n        x=self.relu(x)\n        x=self.conv2d_3(x)\n        print(x.shape)\n        x=self.relu(x)\n        x=self.res_1(x)\n        print(x.shape)\n        x=self.res_2(x)\n        print(x.shape)\n        x=self.res_3(x)\n        print(x.shape)\n        x=self.res_4(x)\n        print(x.shape)\n        x=self.res_5(x)\n        print(x.shape)\n        x=self.deconv2d_1(x)\n        print(x.shape)\n        x=self.relu(x)\n        x=self.deconv2d_2(x)\n        print(x.shape)\n        x=self.relu(x)\n        x=self.deconv2d_3(x)\n        print(x.shape)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Define input shape and batch size"},{"metadata":{"id":"387PVphs8B6Z","trusted":true},"cell_type":"code","source":"input_shape=(256,256,3)\nbatch_size=16","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### creating instance of style model"},{"metadata":{"id":"xuzp3j6ZqAef","trusted":true},"cell_type":"code","source":"style_model = StyleTransferModel()","execution_count":null,"outputs":[]},{"metadata":{"id":"5ycmz_hW-4Du","outputId":"cb83cb6b-d316-43e0-b761-4364af137e75","trusted":true},"cell_type":"code","source":"style_model.print_shape(tf.zeros(shape=(1,*input_shape)))","execution_count":null,"outputs":[]},{"metadata":{"id":"sJXHqAn6KH7t"},"cell_type":"markdown","source":"## TRAINING UTILITY\n\nThe output from decoder is passed to loss model (VGG) from which we extract features and calculate style loss and content loss whose weighted sum provide perceptual loss. Using loss, gradients are calculated with respect to style model's trainable parameters."},{"metadata":{},"cell_type":"markdown","source":"### initializing optimizer for backpropogation"},{"metadata":{"id":"ZIRXbqYuqFGx","trusted":true},"cell_type":"code","source":"optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)","execution_count":null,"outputs":[]},{"metadata":{"id":"UZ9ed5fgGfRa","trusted":true},"cell_type":"code","source":"def train_step(dataset,style_activations,steps_per_epoch,style_model,loss_model,optimizer,checkpoint_path=\"./\",content_weight=1e4,style_weight=1e-2,total_variation_weight=0.004,content_layers_weights=[1],style_layers_weights=[1]*5):\n    batch_losses=[]\n    steps=1\n    save_path=os.path.join(checkpoint_path,f\"model_checkpoint.ckpt\")\n    print(\"Model Checkpoint Path: \",save_path)\n    for input_image_batch in dataset:\n        if steps-1 >= steps_per_epoch:\n            break\n        with tf.GradientTape() as tape:\n            outputs=style_model(input_image_batch)\n            outputs=tf.clip_by_value(outputs, 0, 255)\n            pred_activations=loss_model.get_activations(outputs/255.0)\n            content_activations=loss_model.get_activations(input_image_batch)[\"content\"] \n            curr_loss=preceptual_loss(pred_activations,content_activations,style_activations,content_weight,\n                                      style_weight,content_layers_weights,style_layers_weights)\n            curr_loss += total_variation_weight*tf.image.total_variation(outputs)\n        batch_losses.append(curr_loss)\n        grad = tape.gradient(curr_loss,style_model.trainable_variables)\n        optimizer.apply_gradients(zip(grad,style_model.trainable_variables))\n        if steps%100==0:\n            print(\"checkpoint saved \",end=\" \")\n            style_model.save_weights(save_path)\n            print(f\"Loss: {tf.reduce_mean(batch_losses).numpy()}\")\n        steps+=1\n    return tf.reduce_mean(batch_losses)","execution_count":null,"outputs":[]},{"metadata":{"id":"8nRdNH_tKQUX"},"cell_type":"markdown","source":"## SETTING UP DATASET\n\ntf.data api from tensorflow is used to set up training data."},{"metadata":{"id":"xYDiO-19fdaO","trusted":true},"cell_type":"code","source":"input_path = \"../input/gan-getting-started/photo_jpg\"\nstyle_path = \"../input/gan-getting-started/monet_jpg\"","execution_count":null,"outputs":[]},{"metadata":{"id":"hGMI5nPhUDdc","trusted":true},"cell_type":"code","source":"class TensorflowDatasetLoader:\n    def __init__(self,dataset_path,batch_size=4, image_size=(256, 256),num_images=None):\n        images_paths = [str(path) for path in Path(dataset_path).glob(\"*.jpg\")]\n        self.length=len(images_paths)\n        if num_images is not None:\n            images_paths = images_paths[0:num_images]\n        dataset = tf.data.Dataset.from_tensor_slices(images_paths).map(\n            lambda path: self.load_tf_image(path, dim=image_size),\n            num_parallel_calls=tf.data.experimental.AUTOTUNE,\n        )\n        dataset = dataset.batch(batch_size,drop_remainder=True)\n        dataset = dataset.repeat()\n        dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)\n        self.dataset=dataset\n    def __len__(self):\n        return self.length\n    def load_tf_image(self,image_path,dim):\n        image = tf.io.read_file(image_path)\n        image = tf.image.decode_jpeg(image, channels=3)\n        image= tf.image.resize(image,dim)\n        image= image/255.0\n        image = tf.image.convert_image_dtype(image, tf.float32)\n        return image","execution_count":null,"outputs":[]},{"metadata":{"id":"6ViGmoQOUIIX","trusted":true},"cell_type":"code","source":"loader=TensorflowDatasetLoader(input_path,batch_size=batch_size)","execution_count":null,"outputs":[]},{"metadata":{"id":"xY9j1oX7VH4j","outputId":"0a76c8e2-35e1-47be-abc6-cf93eccb5af8","trusted":true},"cell_type":"code","source":"loader.dataset.element_spec","execution_count":null,"outputs":[]},{"metadata":{"id":"JbKOlY5kVNQw","outputId":"d48ca4a2-c68f-4553-fa79-08602ca03a77","trusted":true},"cell_type":"code","source":"plot_images_grid(next(iter(loader.dataset.take(1))),num_rows=4)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### setting up our style image for training"},{"metadata":{"id":"Nmlu5PagLNEl","trusted":true},"cell_type":"code","source":"# setting up style image\n\nstyle_image_path =os.path.join(style_path,\"05b493ff42.jpg\")\nstyle_image=load_image(style_image_path,dim=(input_shape[0],input_shape[1]),resize=True)\nstyle_image=style_image/255.0","execution_count":null,"outputs":[]},{"metadata":{"id":"v7BfOzJ4EUmv","outputId":"991c6ee6-f3ef-4bf9-dd8e-9bc6dc62fcb5","trusted":true},"cell_type":"code","source":"show_image(style_image)","execution_count":null,"outputs":[]},{"metadata":{"id":"QSLSvYFRVbjv","trusted":true},"cell_type":"code","source":"style_image=style_image.astype(np.float32)\nstyle_image_batch=np.repeat([style_image],batch_size,axis=0)\nstyle_activations=loss_model.get_activations(style_image_batch)[\"style\"]","execution_count":null,"outputs":[]},{"metadata":{"id":"Q9jojrf9KLxt"},"cell_type":"markdown","source":"## TRAINING STYLE MODEL"},{"metadata":{"id":"O-SDoxCf_RpQ","trusted":true},"cell_type":"code","source":"epochs=10\ncontent_weight=1e1\nstyle_weight=1e2\ntotal_variation_weight=0.004","execution_count":null,"outputs":[]},{"metadata":{"id":"MDpGE67EzY2q","outputId":"d6059eb7-7281-4278-8d17-a40759c47314","trusted":true},"cell_type":"code","source":"num_images=len(loader)\nsteps_per_epochs=num_images//batch_size\nprint(steps_per_epochs)","execution_count":null,"outputs":[]},{"metadata":{"id":"UZpRTmkFrkKo","outputId":"13f3cd1f-66d0-48e8-8a6a-e32845219448","trusted":true},"cell_type":"code","source":"model_save_path=\"model_checkpoint\"","execution_count":null,"outputs":[]},{"metadata":{"id":"917zhP1TwvDO","trusted":true},"cell_type":"code","source":"os.makedirs(model_save_path,exist_ok=True)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### enabling mix precision and jit for training optimization"},{"metadata":{"id":"NF16ky1uWDmJ","trusted":true},"cell_type":"code","source":"try:\n    policy = tf.keras.mixed_precision.experimental.Policy('mixed_float16')\n    tf.keras.mixed_precision.experimental.set_policy(policy) \nexcept:\n    pass","execution_count":null,"outputs":[]},{"metadata":{"id":"Kgk-Om0pWEhg","trusted":true},"cell_type":"code","source":"try:\n    tf.config.optimizer.set_jit(True)\nexcept:\n    pass","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Loading previous saved checkpoints if exists"},{"metadata":{"id":"XeXbREFZWKHT","outputId":"7283f071-44f5-4345-eba9-6f1f454e0b74","trusted":true},"cell_type":"code","source":"if os.path.isfile(os.path.join(model_save_path,\"model_checkpoint.ckpt.index\")):\n    style_model.load_weights(os.path.join(model_save_path,\"model_checkpoint.ckpt\"))\n    print(\"resuming training ...\")\nelse:\n    print(\"training scratch ...\")","execution_count":null,"outputs":[]},{"metadata":{"id":"vINKbqM5I-sG","outputId":"9c7990f5-f061-434f-fb09-41cf0698de4a","trusted":true},"cell_type":"code","source":"epoch_losses=[]\nfor epoch in range(1,epochs+1):\n    print(f\"epoch: {epoch}\")\n    batch_loss=train_step(loader.dataset,style_activations,steps_per_epochs,style_model,loss_model,optimizer,\n                          model_save_path,\n                          content_weight,style_weight,total_variation_weight,\n                          content_layers_weights,style_layers_weights)\n    style_model.save_weights(os.path.join(model_save_path,\"model_checkpoint.ckpt\"))\n    print(\"Model Checkpointed at: \",os.path.join(model_save_path,\"model_checkpoint.ckpt\"))\n    print(f\"loss: {batch_loss.numpy()}\")\n    epoch_losses.append(batch_loss)","execution_count":null,"outputs":[]},{"metadata":{"id":"w4lvscOlMezg","outputId":"a0bbfcbf-e846-4917-ea0c-9fa7d59cd28a","trusted":true},"cell_type":"code","source":"plt.plot(epoch_losses)\nplt.xlabel(\"Epochs\")\nplt.ylabel(\"Loss\")\nplt.title(\"Training Process\")\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"id":"Pq_JZOUDg2Hu","outputId":"9577231c-a6f5-4b59-a7d3-9a76e025b263","trusted":true},"cell_type":"code","source":"if os.path.isfile(os.path.join(model_save_path,\"model_checkpoint.ckpt.index\")):\n    style_model.load_weights(os.path.join(model_save_path,\"model_checkpoint.ckpt\"))\n    print(\"loading weights ...\")\nelse:\n    print(\"no weights found ...\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## GENERATING ARTWORK"},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_images_grid(next(iter(loader.dataset.take(1))),num_rows = 4)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"styled_images = []\n\nfor images in loader.dataset.take(1):\n    images = images * 255.0\n    images= tf.image.resize(images,(1024,1024))\n    generated_images = style_model(images)\n    generated_images = np.clip(generated_images,0,255)\n    generated_images = tf.image.resize(generated_images,(256,256)).numpy()\n    generated_images = generated_images.astype(np.uint8)\n    for image in generated_images:\n        styled_images.append(image)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_images_grid(styled_images,num_rows = 4)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## ZIPPING SUBMISSION DATA"},{"metadata":{"trusted":true},"cell_type":"code","source":"os.makedirs(\"images\",exist_ok=True)","execution_count":null,"outputs":[]},{"metadata":{"id":"ge0gD78YTX63","trusted":true},"cell_type":"code","source":"i = 1\nfor images in loader.dataset.take(500):\n    images = images * 255.0\n    images= tf.image.resize(images,(1024,1024))\n    generated_images = style_model(images)\n    generated_images = np.clip(generated_images,0,255)\n    generated_images = tf.image.resize(generated_images,(256,256)).numpy()\n    generated_images = generated_images.astype(np.uint8)\n    for image in generated_images:\n        img = Image.fromarray(image)\n        img.save(os.path.join(\"images\",f\"{i}.jpg\"))\n        i=i+1","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import shutil\n\nshutil.make_archive('/kaggle/working/images/', 'zip', 'images')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"shutil.make_archive('/kaggle/working/model_checkpoint/', 'zip', 'model_checkpoint')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!rm -rf images\n!rm -rf model_checkpoint","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}