# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# + [markdown] id="nq_B010XthHd" colab_type="text"
#
# #<center> <b>Neural Transfer Using PyTorch</b></center>
# ## <center> Studying Gradient Descent Methods</center>
# ###<center><EMAIL></center>
#
#
# This is a quick extension of the PyTorch tutorial by <NAME> [NST Tutorial](https://pytorch.org/tutorials/advanced/neural_style_tutorial.html), but designed to run within the Colaboratory environment with a Google Drive as the permanent storage area.  The particular choice for using Google Drive is driven by my personal transition to a Chromebox/book environment; therefore, optimization of this change requires minimizing any transfers to/from my local drive.
#
# ##Introduction
# There is a lot of activity in this notebook since we need to:
# 1.  Demostrate moving data to/from a Google Drive
# 2.  Demonstrate installing and using PyTorch
# 3.  Using gradient methods in PyTorch while leveraging and extending a previously trained model.
#
# My main focus is on the gradient methods, but the other two pieces are also informative for solving future problems.
# There are a couple of challenges, namely:
#
# This notebook doesn't spend a lot of time on the Neural Style Transfer algorithm (since it is readily available at [Tutorial ](https://pytorch.org/tutorials/advanced/neural_style_tutorial.html#sphx-glr-download-advanced-neural-style-tutorial-py)), rather, the focus is on leveraging gradient descent and Colaboratory to do something interesting.  This notebook has four main paragraphs:
#
# 1.  Setting up the  Environment
# 2.  Performing the Initial Image Processing
# 3.  Using Gradient Descent for Neural Transfer
# 4.  Saving the Result
#
# NOTE:  This can be done using just a cpu; however, the GPU processing is much faster.  Consequently using a GPU is recommended for this notebook.  Since this notebook is assumed to be running in Colaboratory, you should change the runtime to include a GPU.

# + [markdown] id="aWEi3nSy56Kb" colab_type="text"
# ##1.  Set up the Environment
#
# After we load the torch capabilities, we need to restart to runtime so that the PIL capabilities are synched between the default on the Colaboratory environment (the old verion 4.x) and the requirement for torch (the newer verion 5.x).  Otherwise everything is fairly standard, so the main steps are:
#
# <ol type="a">
#   <li>Import my standard python environment: os, sys, numpy, scipy, matplotlib, skimage tools</li>
#   <li>Install/import torch and torchvision</li>
#   <li>Update the Python Environment</li>
#   <li>Copy the two working files from my Google Drive</li>
# </ol>
#
#

# + [markdown] id="X3KWNihrsy7X" colab_type="text"
# ###1.a  Import Standard Python

# + [markdown] id="uknRpA3o9yu9" colab_type="text"
# python is a rich environment, but these import the standard imports in the first compute cell and what I consider the specific ones for this notebook in the second cell.

# + id="xCQpN8kotKWJ" colab_type="code" colab={}
#General python includes
#skimage is use for much of the preprocessing
import os, sys
import os.path
import numpy as np
import scipy as sp
from pprint import pprint, pformat
from itertools import product
import skimage
from skimage import io, transform
#
#The plot capabilities
import matplotlib
import matplotlib.image as mpimg
import matplotlib.pyplot as plt
#Since these are images, turn off the grids
matplotlib.rc('axes',**{'grid':False})
#Output system specific informaiton
print("Current Working Directory: {:s}".format(os.getcwd()))
_content_path_ = os.getcwd()

# + id="AnHg1ik9t0Y1" colab_type="code" colab={}
#To perform deep copies
import copy

# + [markdown] id="uQaVXgbfuQl4" colab_type="text"
# ###1.b  Install and Import Torch/Torchvision
#
# The Colaboratory environment is very nice because you have the option of using a GPU.  Once a data scientist moves beyond the simple examples, the choice of using a GPU really isn't a choice; therefore, a user must choose an environment.  The obvious choice is Tensorflow <i>(since it is embedded in Colaboratory)</i>, but personal preferences come in to play choosing between Tensorflow and PyTorch.  I have based my own choice of Torch <i>(I don't use it for my building class labs because it requires a few extra steps.)</i>  on the following observations -
# <ol type="i">
#   <li>PyTorch feels more pythonic</li>
#   <li>PyTorch recomputes the gradient dynamically; therefore, we can easily mix/match gradient techniques within a module</li>
#   <li>The entendability of pretrained Deep Learning models seems more natural</li>
# </ol>
# These are just personal preferences, but they work in this notebook.

# + id="MeFjQNPy6Xxk" colab_type="code" colab={}
#
#Install Torch and TorchVision, torchvision is the home to vgg models
# !pip3 install -U torch torchvision
import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import transforms
from torchvision import models
#  import PIL since torchvision uses PIL
import PIL
from PIL import Image
#
#  Output Information
#
has_cuda = torch.cuda.is_available()
current_device = torch.cuda.current_device() if has_cuda else -1
gpu_count = torch.cuda.device_count() if has_cuda else -1
gpu_name = torch.cuda.get_device_name(current_device) if has_cuda else "NA"
print("Number of devices: {:d}".format(gpu_count))
print("Current GPU Number: {:d}".format(current_device))
print("GPU Name: {:s}".format(gpu_name))
#Set the accelerator variable
accelerator = 'cuda' if has_cuda else 'cpu'
print("Accelerator: {:s}".format(accelerator))

#

# + [markdown] id="2JjjSL-uzKCJ" colab_type="text"
# ###Update the Python Environment

# + [markdown] id="5iUXguzD-gBD" colab_type="text"
# At this point, I apologize for the inconvenience, but restart the current runtime.  After restart, re-execute the previous compute cells and skip this action.  <i>(If anyone knows an easier way to accomplish these updates, please drop me a note.)</i>

# + [markdown] id="iY97_nc--aVy" colab_type="text"
# ###1.d  Copy the Two Working Files from Google Drive
#
# Per my transition to using the web for data storage, this notebook assumes all permanent storage is on a Google Drive.  Anyone using Colaboratory has a Google Drive; therefore, this is a nice optimization.  To optimize this approach, a personal library was developed <i>(and available at [Connection To Google Drive](https://github.com/drdavidrace/colab_gdrive))</i> that is more "Linux" like to avoid obtaining the file_id for each individual file of interest.  The library uses the "file path" to identify a file and contains methods for copying the file to the local directory.  For this we need to :
#
# *   Connect to the Google Drive
# *   Copy the files to the local VM
#
# We will use the connection to save the file at a later point.
#
# Connect to your Google Drive -

# + id="M-fKEDvL-xCz" colab_type="code" colab={}
#Connect to your Google Drive
import logging
#
# !pip uninstall --yes colab_gdrive
# !pip install -U -q git+https://github.com/drdavidrace/colab_gdrive.git
# !pip list | grep -i colab
# !pip install -U -q PyDrive
#
#  Connect to a ColabGDrive
from colab_gdrive import colab_gdrive

myGdrive = colab_gdrive.ColabGDrive(logging_level=logging.ERROR)
pprint(myGdrive.is_connected())
pprint(myGdrive.getcwd())

# + [markdown] id="8mZQcTsp_khf" colab_type="text"
# Download your images -
#
# This uses a form, which must be run before proceeding.  

# + id="u2Q3-i1E5uK5" colab_type="code" colab={}
#@title Run this cell to update the files and default locations
content_photo = 'Self.jpg' #@param {type:"string"}
style_photo = 'picasso.jpg' #@param {type:"string"}
google_drive_path = 'BigDataTraining/FunWithGradientDescent' #@param{type:"string"}

# + id="JEiHz1XH_nKv" colab_type="code" colab={}
#Define files to copy 
#change the next three lines to match your files
content_photo = content_photo.strip()
style_photo = style_photo.strip()
google_drive_path = google_drive_path.strip()
#
#  Change the google drive path
#
myGdrive.chdir(google_drive_path)
#
data_files = [content_photo, style_photo]
#Check local file existence
files_to_copy = []
for df in data_files:
  if not os.path.isfile(df):
    image_name = df
    files_to_copy.append(image_name)
# #copy files
pprint(files_to_copy)
if files_to_copy:
  myGdrive.copy_from(files_to_copy)
#Check local file existence
for df in data_files:
  if os.path.isfile(df):
    pprint("Found local version of " + df)
  else:
    pprint("Did not find local version of " + df)

# + [markdown] id="fxtabQiP_1xz" colab_type="text"
# These files should be stored in the "/content" directory, so run a visual check for the files.

# + id="8u4BxN0U_42r" colab_type="code" colab={}
# !pwd
# !ls -alF

# + [markdown] id="iHMGD2dt7iB6" colab_type="text"
# ###  Conclusion

# + [markdown] id="XfHuNdPk7lY_" colab_type="text"
# At this point the environment is set up and the data is ready for processing.

# + [markdown] id="l4vIPBTgAOST" colab_type="text"
# ##2.  Perform the Initial Image Processing

# + [markdown] id="6RV8PWYIAaAL" colab_type="text"
# I have always liked the Picasso effects, but style tranfer alone isn't sufficient to achieve an interesting look.  To achieve the desired effects, the initial image processing performs these other steps:
#
# *  Permuting the squares of the images (I experimented with different size squares, but 3x3 decompositon seems to work well for my photo.)
# *  Rotating the squares of the images (I don't rotate many of the squares, but this gives a nice look.)
# *  Doing a blue color shift (this can be adjusted.)
#
# Since I am using a 224x224 image (primarily for my personal photo), I first sample to 225x225 for the rotations.  Then later resample to the target size of 224x224.
#
# For the permutation of the square, this supports both a random permutation or a fixed permutation.  Random was nice to start, but I eventually went with fixed since I am only working with a single image.
#
# For the rotations, this codel only supports rotations of 0, $\frac{\pi}{2}$,  $\pi$, and $\frac{\pi}{2}$ degrees (so I don't have to do some type of mapping from the other rotated squares to the target image squares).  
#
# For the color shift, this is rather rudimentary.  This just adds a shift of <i>blue_shift</i> to the blue and subtracts $\frac{blue_shift}{2}$ from the other two colors. 
#
# ><i>Note:  After working with this a little, I have chosen to not do the rotation.  I have left the code so other can have fun with it as desired.</i>
#
#
#
# Set the processing variables:  (You must run the form cell before the values will be set.)

# + id="uia68stpAdXX" colab_type="code" colab={} cellView="form"
#@title  After running this cell the first time, it changes automatically on a change.
do_permute = True #@param ["False", "True"]{type:"raw"}
do_rotation = True #@param ["False", "True"]{type:"raw"}
do_shift = True #@param ["False", "True"]{type:"raw"}
random_permute = False #@param ["False", "True"] {type:"raw"}
blue_shift = .1 #@param {type:"number"}
target_image_rows = 224 #@param {type:"integer"}
target_image_cols = 224 #@param {type: "integer"}
number_blocks = 3 #@param {type: "integer"}
content_image_file = "Self.jpg" #@param{type: "string"}
style_image_file = "picasso.jpg" #@param{type: "string"}
stage_one_output_file = "working_image_one.jpg" #@param{type: "string"}

# + [markdown] id="Xe1ls3hUA0xW" colab_type="text"
# Define the basic functions.

# + id="Zn3jdENzA29e" colab_type="code" colab={}
from math import sqrt
#
#  functions
#
def make_content_name(in_name):
  '''
  Purpose:  Create a complete file name based upon the base directory: /content
  
  Input:
    in_name - The name of the input file
    
  Uses:
    _content_path_ - The global variable that stores the current content path
    
  Output:
    The concatenation of the content_path with the in_name
  '''
  return os.path.join(_content_path_,in_name)

def print_move_matrix(in_matrix):
  '''
  Purpose:  Only used during debugging, but left here.
  '''
  m,nx,ny = in_matrix.shape
  for i,j in product(range(nx),range(ny)):
    pprint('{:d},{:d} -> {:d},{:d}]'.format(i,j,in_matrix[0,i,j],in_matrix[1,i,j]))
def compute_start_rc(k, num_sq, sq_size):
  '''
  Purpose:  compute the start row and start pixel of a particular square used in the permuatation of the squares
  
  Inputs:
    k:  number of the chosen square (row major order)
    num_sq:  The number of squares in each row and column
    sq_size:  The number of pixels in each square (sq_size x sq_size)
    
  Outputs:
    row_start:  The row the data to start cut out
    col_start:  The column of the data to start cut out
  '''
  assert k >= 0
  assert num_sq > 0
  assert sq_size > 0
  
  row = int(k / num_sq)
  col = k - row * num_sq
  row_start = row * sq_size
  col_start = col * sq_size
  return row_start, col_start

def copy_image_part(in_image, k, num_sq):
  '''
  Purpose:  Create a copy of an square cut out of an image
  
  Inputs:
    in_image:  The original image
    k:  The square number to cut out
    num_sq:  The number of square to use for image cut outs
    
  Outputs:
    A copy of the data in the square to cut out
  '''
  nx,ny,m = in_image.shape
  sq_size = int(nx/num_sq)
  assert sq_size * num_sq == nx
  assert nx == ny
  row_start, col_start = compute_start_rc(k, num_sq, sq_size)
  temp_array = in_image[row_start:row_start + sq_size, col_start:col_start + sq_size,:].copy()
  return temp_array

def permute_image(in_image, permute_array, num_sq):
  '''
  Purpose:  Permute the square in the image
  
  Inputs:
    in_image: The original image
    permute_array:  An array defining the permutation
    num_sq:  The number of squares in each row and column
    
  Outputs:
    A copy of the image with the squares permuted
  '''
  nx,ny,m = in_image.shape
  w_image = in_image.copy()
  assert nx == ny
  sq_size = int(nx/num_sq)
  assert sq_size  * num_sq == nx
  work_array = np.zeros((sq_size,sq_size,m))
  temp_array = np.zeros((sq_size,sq_size,m))
  touched = np.full((len(permute_array)),False,dtype=bool)
  cur_no = 0
  work_array = w_image[:sq_size,:sq_size,:].copy()
  next_no = 0
  while not np.all(touched):

    next_no = permute_array[cur_no]
    if not touched[next_no]:
      temp_array = copy_image_part(w_image,next_no,num_sq)
      row_start, col_start = compute_start_rc(next_no, num_sq, sq_size)
      w_image[row_start:row_start + sq_size, col_start:col_start + sq_size,:] = work_array.copy()
      work_array = temp_array.copy()
    else:
      for i in range(len(permute_array)):
        if not touched[i]:
          cur_no = i
          break
      work_array = copy_image_part(w_image, cur_no, num_sq)
      next_no = permute_array[cur_no]
      temp_array = copy_image_part(w_image,next_no,num_sq)
      row_start, col_start = compute_start_rc(next_no, num_sq, sq_size)
      w_image[row_start:row_start + sq_size, col_start:col_start + sq_size,:] = work_array.copy()
      work_array = temp_array.copy()
    touched[next_no] = True
    cur_no = next_no
  return w_image


# + [markdown] id="A5mfzzRQBEb7" colab_type="text"
# Process the Content Image with permutations, rotations and color shifts.  <i>(These aren't necessary, but I found the result more interesting.)</i>

# + id="kYUX34nzBHjd" colab_type="code" colab={}
#Set 
assert target_image_rows == target_image_cols, "The target number of rows " + str(target_image_rows) + " must equal the number of columns " + str(target_image_cols)
img_size = target_image_rows

image_path = make_content_name(content_image_file)
num_sq = int(number_blocks)
t_size = 0
if( int(target_image_rows/num_sq) * num_sq == target_image_rows):
  t_size = target_image_rows
else:
  t_size = (int(target_image_rows/num_sq) + 1) * num_sq

print("Target Content Image Size {:d}".format(target_image_rows))
print("Working Image Size: {:d}".format(t_size))

out_path = make_content_name(stage_one_output_file)

# matrix_0, matrix_90, matrix_180, matrix_270 = define_rotation_matrices(t_size, num_sq)


# content_image = 
content_image = io.imread(image_path)
content_array = transform.resize(content_image,(t_size,t_size))
print("Content Image - Resampled")
pprint(content_array.shape)
plt.imshow(content_array)
plt.show()
tot_sq = num_sq * num_sq
#permute image
#set this if you want to use a specific permutation
rand_perm_in = [8, 3, 1, 2, 7, 6, 0, 5, 4]
np.random.seed(0)
image_out = None
if do_permute:
  rand_perm = np.random.permutation(tot_sq) if random_permute else rand_perm_in
  pprint("The permutation")
  pprint(rand_perm)
  image_out = permute_image(content_array,rand_perm,num_sq)
else:
  image_out = content_array.copy()

plt.imshow(image_out)
plt.show()

if do_shift:
  print("Blue Shift {:f}".format(blue_shift))
  image_out[:,:,2] = image_out[:,:,2] + 2.0 * blue_shift
  image_out[:,:,0] = image_out[:,:,0] - blue_shift/3.0
  image_out[:,:,1] = image_out[:,:,1] - 2.0 * blue_shift/3.0
  #clip
  image_out = np.clip(image_out,0.0,1.0)

plt.imshow(image_out)
plt.show()

if do_rotation:
  nx,ny,m = image_out.shape
  sq_size = int(nx/num_sq)
  p_rotate = 1.0/8.
  move_matrix = None
  rot_val = 1
  np.random.seed(np.random.randint(0,high=np.iinfo(np.int32).max))
  for k in range(num_sq * num_sq):
    r_val = np.random.rand()
    row_start,col_start = compute_start_rc(k,num_sq, sq_size)
    rotate = False
    rot_angle = 0.0
    if r_val < 3 * p_rotate:
      rot_val = (rot_val + 1) %3
      rotate = True
    if rotate and (rot_val ==0):
      rot_angle = 90.
      rotate = True
    elif rotate and (rot_val == 1):
      rot_angle = 180.
      rotate = True
    elif rotate and (rot_val == 2):
      rot_angle = 270.
      rotate = True
    if rotate:
      w_array = copy_image_part(image_out, k, num_sq)
      t_array = np.zeros(w_array.shape)
      wx, wy, wm = w_array.shape
      t_array = transform.rotate(w_array,rot_angle)
      image_out[row_start:row_start+sq_size,col_start:col_start+ sq_size,:] = t_array
    
if do_permute:
  permute_array = [0, 3, 2, 7, 4, 5, 6, 1, 8]
  pprint(rand_perm)
  image_out = permute_image(image_out,permute_array,num_sq)
  
show_image = image_out * 255.
show_image = np.array(show_image).astype(int)
show_out = np.uint8(show_image)
plt.imshow(image_out)
plt.show()
# !rm -rf out_path
io.imsave(out_path,show_out)

# + [markdown] id="2NuFe4NGBmdL" colab_type="text"
# Check the files.

# + id="5BVVz6B0Bo4-" colab_type="code" colab={}
# !ls -alF

# + [markdown] id="SsDBKWOvCRlb" colab_type="text"
# ##3.  Use "Gradient Descent" for Neural Style Transfer
#
# In this process we read in the content image and style images to perform Neural Style Transfer.  The differnce in this part of the notebook is the way we use autograd.  The use of autograd within PyTorch is very powerful, namely, there is nothing restricting us to a single gradient descent optimization technique during the optimization process.  The PyTorch autograd method allows us some flexibility to both converge faster to a solution and smooth the solution as we approach the end of the process.

# + [markdown] id="mK2hndnhG3he" colab_type="text"
# ###3.1  Load the Images onto the GPU

# + id="_qyS4sLlCYSx" colab_type="code" colab={}
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
#load the images
content_file = 'working_image_one.jpg'
style_file = 'picasso.jpg'
loader = transforms.Compose([
    transforms.Resize([img_size,img_size],PIL.Image.LANCZOS),  # scale imported image
    transforms.ToTensor()])
#
content_file = make_content_name(content_file)
style_file = make_content_name(style_file)
pprint(content_file)
pprint(style_file)
content_image = Image.open(content_file)
content_image = loader(content_image).to(device,torch.float)
style_image = Image.open(style_file)
style_image = loader(style_image).to(device,torch.float)
#
#  Build the target image
#
target_image = content_image.clone()

pprint(type(content_image))
pprint(content_image.device)
pprint(content_image.shape)
pprint(type(style_image))
pprint(style_image.device)
pprint(style_image.shape)
pprint(type(target_image))
pprint(target_image.device)
pprint(target_image.shape)

# + [markdown] id="OBMqmBllDrL1" colab_type="text"
# Show the images

# + id="eEEu8XVaDtHA" colab_type="code" colab={}
unloader = transforms.ToPILImage()  # reconvert into PIL image

def imshow(tensor, title=None):
    image = tensor.cpu().clone()  # we clone the tensor to not do changes on it
    image = image.squeeze(0)      # remove the fake batch dimension
    image = unloader(image)
    plt.imshow(image)
    if title is not None:
        plt.title(title)
    plt.pause(0.001) 
    
plt.figure()
imshow(style_image, title='Style Image')

plt.figure()
imshow(content_image, title='Content Image')

plt.figure()
imshow(target_image, title='Target Image')


# + [markdown] id="1tYJOK0RHKQM" colab_type="text"
# ###3.2 Define the Functions that Generate a "Loss" Function

# + [markdown] id="MlDfm80AthHu" colab_type="text"
# This isn't a true PyTorch Loss function, but rather a measure of the way we want to converge the image between the content and style.  Autograd can work with this looser definition of a loss as easily as it works with an Autograd Loss function.
#
#

# + [markdown] id="Q3ia8CwuthHx" colab_type="text"
# Loss Functions
# --------------
# Define the loss functions for the style and content
#
#
#

# + id="hyqsUTKethHz" colab_type="code" colab={}
class ContentLoss(nn.Module):

    def __init__(self, target,):
        super(ContentLoss, self).__init__()
        # we 'detach' the target content from the tree used
        # to dynamically compute the gradient: this is a stated value,
        # not a variable. Otherwise the forward method of the criterion
        # will throw an error.
        self.target = target.detach()

    def forward(self, input):
        self.loss = F.mse_loss(input, self.target)
        return input
      
def gram_matrix(input):
    a, b, c, d = input.size()  # a=batch size(=1)
    # b=number of feature maps
    # (c,d)=dimensions of a f. map (N=c*d)

    features = input.view(a * b, c * d)  # resise F_XL into \hat F_XL

    G = torch.mm(features, features.t())  # compute the gram product

    # we 'normalize' the values of the gram matrix
    # by dividing by the number of element in each feature maps.
    return G.div(a * b * c * d)
  
class StyleLoss(nn.Module):

    def __init__(self, target_feature):
        super(StyleLoss, self).__init__()
        self.target = gram_matrix(target_feature).detach()

    def forward(self, input):
        G = gram_matrix(input)
        self.loss = F.mse_loss(G, self.target)
        return input
      
def content_and_style(style_losses, content_losses, style_weight, content_weight):
  style_score = 0.0
  content_score = 0

  for sl in style_losses:
      style_score += sl.loss
  for cl in content_losses:
      content_score += cl.loss

  style_score *= style_weight
  content_score *= content_weight
  return content_score, style_score


# + [markdown] id="E1oejX4xthID" colab_type="text"
# ## Building the Model
#
#
# Now we need to import a pre-trained neural network. We will use a 19
# layer VGG network like the one used in the paper.
#
# PyTorch’s implementation of VGG is a module divided into two child
# ``Sequential`` modules: ``features`` (containing convolution and pooling layers),
# and ``classifier`` (containing fully connected layers). We will use the
# ``features`` module because we need the output of the individual
# convolution layers to measure content and style loss. Some layers have
# different behavior during training than evaluation, so we must set the
# network to evaluation mode using ``.eval()``.
#
#
#

# + id="GFxIJ0jFthIF" colab_type="code" colab={}
#Global normalization based upon pretrained model
cnn_normalization_mean = torch.tensor([0.485, 0.456, 0.406]).to(device)
cnn_normalization_std = torch.tensor([0.229, 0.224, 0.225]).to(device)
# cnn_normalization_mean = torch.tensor([0.406, 0.456, 0.485]).to(device)
# cnn_normalization_std = torch.tensor([0.225, 0.224, 0.229]).to(device)
#Load the model
cnn = models.vgg19(pretrained=True).features.to(device).eval()
# create a module to normalize input image so we can easily put it in a
# nn.Sequential
class Normalization(nn.Module):
    def __init__(self, mean, std):
        super(Normalization, self).__init__()
        # .view the mean and std to make them [C x 1 x 1] so that they can
        # directly work with image Tensor of shape [B x C x H x W].
        # B is batch size. C is number of channels. H is height and W is width.
        self.mean = torch.tensor(mean).view(-1, 1, 1)
        self.std = torch.tensor(std).view(-1, 1, 1)

    def forward(self, img):
        # normalize img
        return (img - self.mean) / self.std


# + [markdown] id="Sybl7iCythIO" colab_type="text"
# A ``Sequential`` module contains an ordered list of child modules. For
# instance, ``vgg19.features`` contains a sequence (Conv2d, ReLU, MaxPool2d,
# Conv2d, ReLU…) aligned in the right order of depth. We need to add our
# content loss and style loss layers immediately after the convolution
# layer they are detecting. To do this we must create a new ``Sequential``
# module that has content loss and style loss modules correctly inserted.
#
#
#

# + id="zTehFhYqthIP" colab_type="code" colab={}
# desired depth layers to compute style/content losses :
content_layers_default = ['conv2_3_2']
style_layers_default = [ 'conv2_3_1', 'conv2_4_2', 'conv2_5_3']

def get_style_model_and_losses(in_cnn, norm_mean, norm_std,
                               style_img, content_img,
                               content_layers=content_layers_default,
                               style_layers=style_layers_default):
  
  cnn = copy.deepcopy(in_cnn)
  normalization = Normalization(norm_mean, norm_std).to(accelerator)
  #Define list for content_losses and style_losses
  content_losses = []
  style_losses = []
  #Create new model
  model = nn.Sequential(normalization)
  name = ''
  l = 1
  conv_sub_l = 1
  relu_sub_l = 1
  for layer in cnn.children():
    if isinstance(layer, nn.Conv2d):
      name = 'conv2_{:d}_{:d}'.format(l,conv_sub_l)
      conv_sub_l += 1
    elif isinstance(layer, nn.ReLU):
      name = 'relu_{:d}_{:d}'.format(l, relu_sub_l)
      layer = nn.ReLU(inplace=False)
      relu_sub_l += 1
    elif isinstance(layer, nn.MaxPool2d):
      name = 'pool_{:d}'.format(l)
      l += 1
      conv_sub_l = 1
      relu_sub_l = 1
    else:
      raise RuntimeError('Unknown Layer: {}'.format(layer.__class__.__name__))
    #
    model.add_module(name, layer)
    #
    if name in content_layers:
      target = model(content_img).detach()
      content_loss = ContentLoss(target)
      model.add_module("content_loss_{:s}".format(name), content_loss)
      content_losses.append(content_loss)
    #
    if name in style_layers:
      # add style loss:
      target_feature = model(style_img).detach()
      style_loss = StyleLoss(target_feature)
      model.add_module("style_loss_{:s}".format(name), style_loss)
      style_losses.append(style_loss)

#   for layer in model.named_children():
#     pprint(layer[0])


  return model, style_losses, content_losses


# + [markdown] id="1UJMxdGVthIW" colab_type="text"
# ## Define Gradient Descent
#
# This uses the Adam descent algorithm
#
#
#

# + id="0MMhePYXthIa" colab_type="code" colab={}
def get_input_optimizer0(input_img):
    # this line to show that input is a parameter that requires a gradient
    optimizer = optim.Adam([input_img.requires_grad_()], lr=.01)
    return optimizer
  
def get_input_optimizer1(input_img):
    # this line to show that input is a parameter that requires a gradient
    optimizer = optim.Adam([input_img.requires_grad_()], lr=.005)
    return optimizer
  
def get_input_optimizer2(input_img):
    # this line to show that input is a parameter that requires a gradient
    optimizer = optim.Adam([input_img.requires_grad_()], lr=.001)
    return optimizer


# + [markdown] id="2O6OkpetthIe" colab_type="text"
# Finally, we must define a function that performs the neural transfer. For
# each iteration of the networks, it is fed an updated input and computes
# new losses. We will run the ``backward`` methods of each loss module to
# dynamicaly compute their gradients. The optimizer requires a “closure”
# function, which reevaluates the modul and returns the loss.
#
# We still have one final constraint to address. The network may try to
# optimize the input with values that exceed the 0 to 1 tensor range for
# the image. We can address this by correcting the input values to be
# between 0 to 1 each time the network is run.
#
#
#

# + id="ixbCIPllthIg" colab_type="code" colab={}
global last_loss, cur_opt, run_step, cur_rel_err
last_loss = 0.0
cur_opt = None
cur_rel_err = 1.0
run_step = -1
def run_style_transfer(cnn, normalization_mean, normalization_std,
                       content_img, style_img, input_img, num_steps=10000,
                       style_weight=1000000, content_weight=1):
  global last_loss, cur_opt, run_step, cur_rel_err
  """Run the style transfer."""
  print('Building the style transfer model..')
  model, style_losses, content_losses = get_style_model_and_losses(cnn,
                                                                   normalization_mean, normalization_std, style_img, content_img)
  cur_opt = get_input_optimizer1(input_img)
  print(cur_opt)
  print(type(cur_opt))
  cur_rel_err = 0.5
  print('Optimizing..')
  run_step = 0
  while run_step <= num_steps:
    def closure():
      global last_loss, cur_opt, run_step, cur_rel_err
      # correct the values of updated input image
      input_img.data.clamp_(0, 1)
      cur_opt.zero_grad()
      model(input_img)
      content_score, style_score = content_and_style(style_losses, content_losses, style_weight, content_weight)
      loss = content_score + style_score
      loss.backward()
      #This is included to demonstrate changing the gradient algorithm
      if run_step % 50 == 0:
        curr_loss = loss.item()
        cur_rel_err = np.abs((last_loss - curr_loss))
        last_loss = curr_loss
      if run_step % 500 == 0:
        print("step {:8d}:".format(run_step))
        print('Style Loss : {:4f} Content Loss: {:4f}'.format(
            style_score.item(), content_score.item()))
        print()
        plt.figure()
        imshow(input_img, title='Output Image')
        plt.show
      return style_score + content_score
    #
    cur_opt.step(closure)
    if run_step % 1000 == 0:
      if cur_rel_err > 2.0:
        cur_opt = get_input_optimizer2(input_img)
      elif cur_rel_err > 1.0:
        cur_opt = get_input_optimizer1(input_img)
      else:
        cur_opt = get_input_optimizer0(input_img)
      pprint(cur_opt)
    run_step = run_step + 1

  # a last correction...
  input_img.data.clamp_(0, 1)

  return input_img, run_step


# + [markdown] id="epzkedVfthIk" colab_type="text"
# Finally, we can run the algorithm.
#
#
#

# + id="oP8idZVzthIm" colab_type="code" colab={}
content_image_s = torch.stack([content_image])
style_image_s = torch.stack([style_image])
target_image_s = content_image_s.clone()
output, N = run_style_transfer(cnn, cnn_normalization_mean, cnn_normalization_std,
                            content_image_s, style_image_s, target_image_s)
print("Number of steps: {}".format(N))
plt.figure()
imshow(output, title='Output Image')

# sphinx_gallery_thumbnail_number = 4
plt.ioff()
plt.show()


# + [markdown] id="t7LfSQ--KtQT" colab_type="text"
# ##4.  Save the Results
#
# This is an easy step when using the colab_gdrive library.  We save the file locally, then transfer the file to the Google Drive.

# + [markdown] id="khpXt0hRGwpj" colab_type="text"
# Copy the file to local storage

# + id="KP1tKlpzET6Q" colab_type="code" colab={}
#
#  Save file locally
#
out_jpg_file = "style_converged.jpg"
out_jpg_full = make_content_name(out_jpg_file)
torchvision.utils.save_image(output.squeeze(),out_jpg_full)
# !ls -alF
#
#  show the file
#
simple_loader = transforms.Compose([
    transforms.ToTensor()])
#
converge_image = Image.open(out_jpg_full)
converge_image = loader(converge_image).to(device,torch.float)
imshow(converge_image,title="Converged Image")

# + [markdown] id="IbYNkdJKGtm1" colab_type="text"
# Copy the file to the google drive
#
# Recall that we have already set the current working directory with the chdir command, so this is relatively easy.

# + id="qyAaxeIhHdTw" colab_type="code" colab={}
print(myGdrive.getcwd())
print(os.getcwd())
myGdrive.copy_to(out_jpg_file) #This only supports copying from the local cwd

# + [markdown] id="QoOYLEKUSP-e" colab_type="text"
# ## 5.  Conclusions
#
# This notebook demonstrated these important ideas:
# 1.  Using Colaboratory and Google Drive together to solve problems
# 2.  Installing and using Pytorch
# 3.  Using gradient methods within PyTorch while leveraging and extending previously trained models
#
# If you have any questions about the operation or contents of this notebook, please contact <EMAIL>

StarcoderdataJupyterScriptsDedupFiltered