# ---
# 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
# ---

# %reload_ext autoreload
# %autoreload 2

# + [markdown] heading_collapsed=true
# # Flask

# + hidden=true
from wtforms import 

# + hidden=true
# %%writefile /home/justin/projects/nst_flask_app/deployment/templates/index.html

<!DOCTYPE html>
<html>

<head>
  <title>Neural Style Transfer</title>
  <link rel="stylesheet" href="/static/css/main.css">
  <link rel="shortcut icon" href="/static/images/lstm.ico">
  
</head>

<body>
  <div class="container">
    <h1>
      <center>Neural Style Transfer</center>
    </h1>

    {% block content %}
    {% for message in form.style.errors %}
    <div class="flash">{{ message }}</div>
    {% endfor %}

    {% for message in form.content.errors %}
    <div class="flash">{{ message }}</div>
    {% endfor %}

    {% for message in form.steps.errors %}
    <div class="flash">{{ message }}</div>
    {% endfor %}

    <form method=post>

      {{ form.style.label }}
      {{ form.style }}

      {{ form.content.label }}
      {{ form.content }}

      {{ form.steps.label }}
      {{ form.steps }}

      {{ form.submit }}
    </form>
    {% endblock %}

  </div>
</body>

</html>

# + hidden=true
# %%writefile /home/justin/projects/nst_flask_app/deployment/nst_app_flask.py

from flask import Flask, render_template, request
from wtforms import (Form, FileField, IntegerField, SubmitField, validators)


app = Flask(__name__)

class ReusableForm(Form):
    """User entry form for entering specifics for generation"""
    # Style
    style = FileField("Upload the Style Image", validators=[
                     validators.InputRequired()])
    
    # Content
    content = FileField("Upload the Content Image", validators=[
                     validators.InputRequired()])

    # Number of steps
    steps = IntegerField("Enter number of steps",
                         default=110, validators=[validators.InputRequired(),
                                                 validators.NumberRange(min=10, max=200, 
                                                 message='Must be between 10 and 200.')])
    # Submit button
    submit = SubmitField("Transfer Style")


@app.route("/", methods=['GET', 'POST'])
def home():
    """Home page of app with form"""
    
    # Create form
    form = ReusableForm(request.form)
    
    # On form entry and all conditions met
    if request.method == 'POST' and form.validate():
        # Extract information
        content_img = request.form['content']
        style_img = request.form['style']
        steps = int(request.form['steps'])

    # Send template information to index.html
    return render_template('index.html', form=form)

app.run(host='0.0.0.0', port=3224)

# + hidden=true

# -

# # Starlette

# ### If I screw something up, go to templates/index_form_working.html
# ### js/client_working.js
# ### nst_star_app_working.py

# ls statics/css/

# + code_folding=[0]
# %%writefile statics/css/style.css
body {
    background-color: #fff;
}

.no-display {
    display: none;
}

.center {
    margin: auto;
    padding: 10px;
    padding-left:50px;
    padding-right:50px;
    text-align: center;
    font-size: 14px;
}

.title {
    font-size: 30px;
    margin-top: 1em;
    margin-bottom: 1em;
    color: #262626;
}

.content {
    margin-top: 10em;
}

.analyze {
    margin-top: 5em;
}

.upload-label {
    padding: 10px;
    font-size: 12px;
}

.result-label {
    margin-top: 0.5em;
    padding: 10px;
    font-size: 13px;
}

button.choose-file-button {
  width: 200px;
  height: 40px;
  border-radius: 2px;
  background-color: #ffffff;
  border: solid 1px #7052CB;
  font-size: 13px;
  color: #7052CB;
}

button.analyze-button {
    width: 200px;
    height: 40px;
    border: solid 1px #7052CB;
    border-radius: 2px;
    background-color: #7052CB;
    font-size: 13px;
    color: #ffffff;
}

button:focus {outline:0;}

# + code_folding=[2, 15, 24, 27, 32, 41, 44]
# %%writefile templates/index_form.html
<html lang='en'>
    <head>
        <meta charset='utf-8'>
        <link rel='stylesheet' href='/statics/css/style.css'>
        <link rel="shortcut icon" href="/statics/images/favicon.ico" type="image/x-icon">
        <script src='/statics/js/client.js'></script>
    </head>
    <body>
        <div>
            <div class='center'>
                <div class='title'>Neural Style Transfer via optimization with L2 Wasserstein Distance</div>
                <h5>Safari may have some timeout issues. Try Chrome or Firefox</h5>

                    <!-- CONTENT -->
                    <div class='no-display'>
                        <input id='content-input'
                            class='no-display'
                            type='file'
                            name='content'
                            accept='image/*'
                            onchange='showPickedContent(this)'>
                    </div>
                    <button class='choose-file-button' type='button' onclick='showPickerContent()'>Select Content Image</button>
                    <div class='upload-label'>
                        <label id='content-upload-label'>No file chosen</label>
                    </div>
                    <div class='content-img'>
                        <img id='content-image-picked' class='no-display' alt='Content Image' height='200'>
                    </div>
                    
                    <!-- STYLE -->
                    <div class='no-display'>
                        <input id='style-input'
                        class='no-display'
                        type='file'
                        name='style'
                        accept='image/*'
                        onchange='showPickedStyle(this)'>
                    </div>
                    <button class='choose-file-button' type='button' onclick='showPickerStyle()'>Select Style Image</button>
                    <div class='upload-label'>
                        <label id='style-upload-label'>No file chosen</label>
                    </div>
                    <div>
                         <img id='style-image-picked' class='no-display' alt='Style Image' height='200'>
                    </div>

                    <!-- STYLED IMAGE AND ANALYZE -->
                    <!--
                    <h3>Choose number of optimization steps: (150 max)</h3>
                    <input type="number" id = 'steps' value="20">
                    -->
                    
                    <div class='no-display'>
                        <input id='stylized-input'
                        class='no-display'
                        type='file'
                        name='stylized'
                        accept='image/*'
                        onchange='showStylized(this)'>
                    </div>
                    
                    <div class='analyze'>
                        <button id='stylize-button' class='analyze-button' type='button' onclick='stylize()'> 
                        Run Style Transfer</button>
                    </div>
                    <div class='title' id='step-count-display'>
                    </div>
                    <div>
                        <img id='stylized-image' class='no-display' alt='Stylized Image' height='200'>
                    </div>


            </div>
        </div>
    </body>
</html>
# -



# + code_folding=[6]
# %%writefile /home/justin/projects/nst_app/deployment/statics/js/client.js
var el = x => document.getElementById(x);


function showPickerContent(inputId) { el('content-input').click(); }

function showPickedContent(input) {
    el('content-upload-label').innerHTML = input.files[0].name;
    var reader = new FileReader();
    reader.onload = function (e) {
        el('content-image-picked').src = e.target.result;
        el('content-image-picked').className = '';
    }
    reader.readAsDataURL(input.files[0]);
}

function showPickerStyle(inputId) { el('style-input').click(); }

function showPickedStyle(input) {
    el('style-upload-label').innerHTML = input.files[0].name;
    var reader = new FileReader();
    reader.onload = function (e) {
        el('style-image-picked').src = e.target.result;
        el('style-image-picked').className = '';
    }
    reader.readAsDataURL(input.files[0]);
}


function stylize() {
    var contentFile = el('content-input').files;
    if (contentFile.length != 1) {
        alert('Please select 1 content image!');
    } 
    
    var styleFile = el('style-input').files;
    if (styleFile.length != 1) {
        alert('Please select 1 style image!');
    } 
    /*
    var ori_steps = el('steps').value;
    if (ori_steps < 1 || ori_steps > 150) {
        alert('Steps must be between 1 and 150!');
    }
    */
    
    el('stylize-button').innerHTML = 'Styling... Large images may take up to 1 min';
    
    var xhr = new XMLHttpRequest();
    var loc = window.location
   
    xhr.open('POST', `${loc.protocol}//${loc.hostname}:${loc.port}/stylize`, true);
    xhr.onerror = function() {alert (xhr.responseText);}
    
    var fileData = new FormData();
    fileData.append('content', contentFile[0]);
    fileData.append('style', styleFile[0]);
    /*
    figure out how to send steps without corrupting the style image?
    trying fileData.append('steps', steps) corrupts. Number(steps) corrupts
    */
    xhr.send(fileData);
    
    xhr.onload = function(e) {
        if (this.readyState === 4) {
            var response = JSON.parse(e.target.responseText);
            var byte_img = response['stylized_image']
            var height = response['h']
            var width = response['w']
            var steps = response['steps']
            el('stylized-image').style.width = width;
            el('stylized-image').style.height = height;
            el('stylized-image').className = '';
            el('stylized-image').src = "data:image/png;base64," + byte_img;
            el('stylize-button').innerHTML = 'Stylized. Rerun Style Transfer?';
            el('step-count-display').innerHTML = `Stylized for ${steps} steps`;
        }
    }
    
}

/*            
            showStylized()
function showStylized() {
    var reader = new FileReader()
    var file = el('stylized-image')
    reader.readAsDataURL(file.toDataURL)
}
*/
# -



os.path.join(cur_dir, 'statics')

# +
# StaticFiles??

# + code_folding=[]
# %%writefile /home/justin/projects/nst_app/deployment/nst_star_app.py

import vgg_classes as vg
import torch
import matplotlib.pyplot as plt

from starlette.applications import Starlette
from starlette.responses import JSONResponse, HTMLResponse, RedirectResponse, FileResponse
from starlette.staticfiles import StaticFiles
#from wtforms import (Form, FileField, IntegerField, SubmitField, validators)

from pathlib import Path
from io import BytesIO
import sys
import uvicorn
import aiohttp
import asyncio
import os
import json
import requests
import base64 
from PIL import Image as PILImage
import torchvision.transforms as transforms

app = Starlette(debug=True, template_directory='templates')
cur_dir = os.getcwd()
app.mount('/statics', StaticFiles(directory='statics'), name='statics')

# setup
# use gpu
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# choose layers
cont_lyrs = ['relu1_1']
stl_lyrs = ['relu1_1', 'relu1_2', 'relu2_1', 'relu2_2', 'relu3_1'] 
# NN choices
pool_type = 'avg'
cont_loss_type = 'mse'
stl_loss_type = 'l2'

# load in gatys vgg, use as base for our vgg
model_dir = os.getcwd() + '/models/'
vgg_g = vg.VGG_g(pool='avg')
vgg_g.load_state_dict(torch.load(model_dir + 'vgg_conv.pth'))


@app.route("/", methods=["GET"])
async def index(request):
    template = app.get_template('index_form.html')
    content = template.render(request=request)
    return HTMLResponse(content)

@app.route("/stylize", methods=["POST"])
async def stylize(request):
    data = await request.form()
    content = await (data["content"].read())
    style = await (data["style"].read())
    steps = 30
    decoded_bytes_img, h, w = do_nst_stylize(content, style, steps)
    print('sending bytes img in jsonresponse')
    return JSONResponse({'stylized_image': decoded_bytes_img,
                         'h': h,
                         'w': w,
                         'steps': steps})
    
def encode(tensor):
    image = tensor.cpu().clone()
    image = image.squeeze(0)
    image = vg.unloader(image)
    image[image>1] = 1
    image[image<0] = 0
    to_pil = transforms.ToPILImage()
    pil_image = to_pil(image)
#     return pil_image
    buff = BytesIO()
    pil_image.save(buff, format="JPEG")
    return base64.b64encode(buff.getvalue()).decode("utf-8")    

def do_nst_stylize(content, style, steps=10):
    content_img = vg.image_loader(BytesIO(content))
    style_img = vg.image_loader(BytesIO(style))
    synth = content_img.clone()
    bs, c, h,w = content_img.shape
    nst = vg.TransferStyle(base_cnn=vgg_g.to(device),
                           cont_loss_type=cont_loss_type,
                           stl_loss_type=stl_loss_type,
                           cont_img=content_img,
                           stl_img=style_img,
                           pool_type=pool_type,
                           cont_lyrs=cont_lyrs,
                           stl_lyrs=stl_lyrs,
                           device=device,
                           resize_conv=False,
                           downsample=False,
                           approx_stl=True)
    output = nst.run_style_transfer(synth, 'adam', steps, 1e0, style_off=False,
                                    content_off=False, show_step=int(steps/5),
                                    loss_plots=False, show_synth=False)
    o_img_data = encode(output)
    return o_img_data, h, w

if __name__ == "__main__":
    if "serve" in sys.argv:
        port = int(os.environ.get("PORT", 8008)) 
        uvicorn.run(app, host="0.0.0.0", port=port)
# -



# mv statics/images/picasso.ico statics/images/picasso.jpg

# ls statics/images/







# + code_folding=[0]
# # %load /home/justin/projects/nst_app/deployment/nst_star_app_working.py

import vgg_classes as vg
import torch
import matplotlib.pyplot as plt

from starlette.applications import Starlette
from starlette.responses import JSONResponse, HTMLResponse, RedirectResponse, FileResponse
from starlette.staticfiles import StaticFiles
from wtforms import (Form, FileField, IntegerField, SubmitField, validators)

from pathlib import Path
from io import BytesIO
import sys
import uvicorn
import aiohttp
import asyncio
import os
import json
import requests
import base64 
from PIL import Image as PILImage
import torchvision.transforms as transforms

app = Starlette(debug=True, template_directory='templates')
app.mount('/statics', StaticFiles(directory='statics'), name='statics')

# setup
# use gpu
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# choose layers
cont_lyrs = ['relu1_1']
stl_lyrs = ['relu1_1', 'relu1_2', 'relu2_1', 'relu2_2', 'relu3_1'] 
# NN choices
pool_type = 'avg'
cont_loss_type = 'mse'
stl_loss_type = 'l2'

# load in gatys vgg, use as base for our vgg
model_dir = os.getcwd() + '/models/'
vgg_g = vg.VGG_g(pool='avg')
vgg_g.load_state_dict(torch.load(model_dir + 'vgg_conv.pth'))


@app.route("/", methods=["GET"])
async def index(request):
    template = app.get_template('index_form.html')
    content = template.render(request=request)
    return HTMLResponse(content)


@app.route("/upload", methods=["POST", "GET"])# "GET"
async def upload(request):
    if request.method == "POST":
        data = await request.form()
        content = await (data["content"].read())
        style = await (data["style"].read())
    #     steps = await (data["steps"].read())
        steps = 10
        return do_nst(content, style, steps)
    elif request.method == "GET":
        return RedirectResponse("/")

@app.route("/stylize", methods=["POST"])
async def stylize(request):
    data = await request.form()
    content = await (data["content"].read())
    style = await (data["style"].read())
    steps = 10
    decoded_bytes_img = do_nst_stylize(content, style, steps)
    print('sending bytes img in jsonresponse')
    return JSONResponse({'stylized_image': decoded_bytes_img})
#     return RedirectResponse("/")
#     return RedirectResponse("/test")

@app.route("/test")
def test(request):
    return HTMLResponse(
        """
        <h1>A test page?</h1>
        """)
    
def encode(tensor):
    image = tensor.cpu().clone()
    image = image.squeeze(0)
    image = vg.unloader(image)
    image[image>1] = 1
    image[image<0] = 0
    to_pil = transforms.ToPILImage()
    pil_image = to_pil(image)
#     return pil_image
    buff = BytesIO()
    pil_image.save(buff, format="JPEG")
    return base64.b64encode(buff.getvalue()).decode("utf-8")    

def do_nst_stylize(content, style, steps=10):
    content_img = vg.image_loader(BytesIO(content))
    style_img = vg.image_loader(BytesIO(style))
    synth = content_img.clone()
    bs, c, h,w = content_img.shape
    nst = vg.TransferStyle(base_cnn=vgg_g.to(device),
                           cont_loss_type=cont_loss_type,
                           stl_loss_type=stl_loss_type,
                           cont_img=content_img,
                           stl_img=style_img,
                           pool_type=pool_type,
                           cont_lyrs=cont_lyrs,
                           stl_lyrs=stl_lyrs,
                           device=device,
                           resize_conv=False,
                           downsample=False,
                           approx_stl=True)
    output = nst.run_style_transfer(synth, 'adam', steps, 1e0, style_off=False,
                                    content_off=False, show_step=int(steps/5),
                                    loss_plots=False, show_synth=False)
    o_img_data = encode(output)
    return o_img_data

def do_nst(content, style, steps=10):
    content_img = vg.image_loader(BytesIO(content))
    style_img = vg.image_loader(BytesIO(style))
    synth = content_img.clone()
    bs, c, h,w = content_img.shape
    nst = vg.TransferStyle(base_cnn=vgg_g.to(device),
                           cont_loss_type=cont_loss_type,
                           stl_loss_type=stl_loss_type,
                           cont_img=content_img,
                           stl_img=style_img,
                           pool_type=pool_type,
                           cont_lyrs=cont_lyrs,
                           stl_lyrs=stl_lyrs,
                           device=device,
                           resize_conv=False,
                           downsample=False,
                           approx_stl=True)
    output = nst.run_style_transfer(synth, 'adam', steps, 1e0, style_off=False,
                                    content_off=False, show_step=int(steps/5),
                                    loss_plots=False, show_synth=False)
    o_img_data = encode(output)
    return HTMLResponse(
        f"""
        <html>
        <figure class="figure">
          <img src="data:image/png;base64, {o_img_data}" class="figure-img img-thumbnail input-image", height="{h}", width="{w}">
        </figure>        
        </html>
    """)

if __name__ == "__main__":
    if "serve" in sys.argv:
        port = int(os.environ.get("PORT", 8008)) 
        uvicorn.run(app, host="0.0.0.0", port=port)
# -





from starlette.responses import JSONResponse, HTMLResponse, RedirectResponse, FileResponse

# +
# FileResponse??
# -

template = app.get_template('index.html')
content = template.render(request=request, content_image = 'dancing.jpg')

# +
img_dir = './images/<NAME>/'
img_dir2 = './images/'
content_name = 'Scan4.jpg'
style_name = 'scream.jpg'

content_img = vg.image_loader(f'{img_dir}{content_name}')
style_img = vg.image_loader(f'{img_dir2}{style_name}')
# -

nst = vg.TransferStyle(base_cnn=vgg_g.to(device), cont_loss_type=cont_loss_type, stl_loss_type=stl_loss_type, cont_img=content_img, stl_img=style_img, pool_type=pool_type, cont_lyrs=cont_lyrs, stl_lyrs=stl_lyrs, device=device, resize_conv=False, downsample=False, approx_stl=True)

# + code_folding=[0]
# # setup fake images before user inputs
# content_img = torch.randn((1,3,512,512)).cuda()
# style_img = torch.randn((1,3,512,512)).cuda()

synth = content_img.clone()
# synth = (torch.rand(content_img.size())*0).cuda()

steps = 20
output = nst.run_style_transfer(synth, 'adam', steps, 1e0, style_off=False, content_off=False, show_step=int(steps/5), loss_plots=False, show_synth=False, ax_synth=None, fig=None)
# -



# # previous, working version

# + code_folding=[0, 56, 70, 83, 96, 109, 170, 175]
# %%writefile nst-app.py
# # %load nst-app.py
from starlette.applications import Starlette
from starlette.responses import JSONResponse, HTMLResponse, RedirectResponse

# import torch
from pathlib import Path
from io import BytesIO
import sys
import uvicorn
import aiohttp
import asyncio
import os
import json
import requests
import base64 
import PIL
from PIL import Image as PILImage

from fastai.vision import *
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

import torchvision.transforms as transforms
import torchvision.models as models

import copy

# NST stuff
# desired depth layers to compute style/content losses :
content_layers_default = ['conv_4']
style_layers_default = ['conv_1', 'conv_2', 'conv_3', 'conv_4', 'conv_5']
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# desired size of the output image
imsize = 512 if torch.cuda.is_available() else 128  # use small size if no gpu

loader = transforms.Compose([
    transforms.Resize(imsize),  # scale imported image
    transforms.CenterCrop(imsize), # center crop it
    transforms.ToTensor()])  # transform it into a torch tensor

def image_loader(image):
    image = PILImage.fromarray(np.uint8(image2np(image.data)*255))
    # fake batch dimension required to fit network's input dimensions
    image = loader(image).unsqueeze(0)
    # drop the 4th channel if it exists
    if image.shape[1] != 3:
        print('image shape is {0} (batch, channel, HxW)'.format(image.shape))
        print('reshaping channel')
        image = image[:,:3,:,:]
        
    return image.to(device, torch.float)

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
    
# 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
    
def get_style_model_and_losses(cnn, normalization_mean, normalization_std,
                               style_img, content_img,
                               content_layers=content_layers_default,
                               style_layers=style_layers_default):
    cnn = copy.deepcopy(cnn)

    # normalization module
    normalization = Normalization(normalization_mean, normalization_std).to(device)

    # just in order to have an iterable access to or list of content/syle
    # losses
    content_losses = []
    style_losses = []

    # assuming that cnn is a nn.Sequential, so we make a new nn.Sequential
    # to put in modules that are supposed to be activated sequentially
    model = nn.Sequential(normalization)

    i = 0  # increment every time we see a conv
    for layer in cnn.children():
        if isinstance(layer, nn.Conv2d):
            i += 1
            name = 'conv_{}'.format(i)
        elif isinstance(layer, nn.ReLU):
            name = 'relu_{}'.format(i)
            # The in-place version doesn't play very nicely with the ContentLoss
            # and StyleLoss we insert below. So we replace with out-of-place
            # ones here.
            layer = nn.ReLU(inplace=False)
        elif isinstance(layer, nn.MaxPool2d):
            name = 'pool_{}'.format(i)
        elif isinstance(layer, nn.BatchNorm2d):
            name = 'bn_{}'.format(i)
        else:
            raise RuntimeError('Unrecognized layer: {}'.format(layer.__class__.__name__))

        model.add_module(name, layer)

        if name in content_layers:
            # add content loss:
            target = model(content_img).detach()
            content_loss = ContentLoss(target)
            model.add_module("content_loss_{}".format(i), 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_{}".format(i), style_loss)
            style_losses.append(style_loss)

    # now we trim off the layers after the last content and style losses
    for i in range(len(model) - 1, -1, -1):
        if isinstance(model[i], ContentLoss) or isinstance(model[i], StyleLoss):
            break

    model = model[:(i + 1)]

    return model, style_losses, content_losses   

def get_input_optimizer(input_img):
    # this line to show that input is a parameter that requires a gradient
    optimizer = optim.LBFGS([input_img.requires_grad_()], lr=.5)
    return optimizer

def run_style_transfer(cnn, normalization_mean, normalization_std,
                       content_img, style_img, input_img, num_steps=300,
                       style_weight=1, content_weight=1):
    """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)
    optimizer = get_input_optimizer(input_img)

    print('Optimizing..')
    run = [0]
    while run[0] <= num_steps:

        def closure():
            # correct the values of updated input image
            input_img.data.clamp_(0, 1)

            optimizer.zero_grad()
            model(input_img)
            style_score = 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

            loss = style_score + content_score
            loss.backward()

            run[0] += 1
            if run[0] % 50 == 0:
                print("run {}:".format(run))
                print('Style Loss : {:4f} Content Loss: {:4f}'.format(
                    style_score.item(), content_score.item()))
                print()

            return style_score + content_score

        optimizer.step(closure)

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

    return input_img

# Non NST stuff

def encode(img):
    img = (image2np(img.data) * 255).astype('uint8')
    pil_img = PILImage.fromarray(img)
    buff = BytesIO()
    pil_img.save(buff, format="JPEG")
    return base64.b64encode(buff.getvalue()).decode("utf-8")




app = Starlette()

path = Path("/tmp")
cnn = models.vgg19(pretrained=True).features.to(device).eval()
normalization_mean = torch.tensor([0.485, 0.456, 0.406]).to(device)
normalization_std = torch.tensor([0.229, 0.224, 0.225]).to(device)

# classes = ['healthy', 'junk']
# data = ImageDataBunch.single_from_classes(path, classes, tfms=get_transforms(), size=224).normalize(imagenet_stats)
# learner = create_cnn(data, models.resnet50)
# learner.model.load_state_dict(
#     torch.load("model-weights.pth", map_location="cpu")
# )


@app.route("/upload", methods=["POST"])
async def upload(request):
    data = await request.form()
    content = await (data["content"].read())
    style = await (data["style"].read())
    return do_nst(content, style)


def do_nst(content, style):
    content_img = image_loader(open_image(BytesIO(content)))
    style_img = image_loader(open_image(BytesIO(style)))
    input_img = content_img.clone()
    output = run_style_transfer(cnn, normalization_mean, normalization_std,
                                content_img, style_img, input_img, num_steps = 300, style_weight = 100000)
    # get rid of batch dimension
#     content_img  = content_img.squeeze(0)
#     style_ig = style_img.squeeze(0)
    output = output.squeeze(0) 
#     pred_class,pred_idx,outputs = learner.predict(img)
#     formatted_outputs = ["{:.1f}%".format(value) for value in [x * 100 for x in torch.nn.functional.softmax(outputs, dim=0)]]
#     pred_probs = sorted(
#             zip(learner.data.classes, map(str, formatted_outputs)),
#             key=lambda p: p[1],
#             reverse=True
#         )
#     c_img_data = encode(content_img)
#     s_img_data = encode(style_img)
    o_img_data = encode(output)
    return HTMLResponse(
        """
        <html>
        <figure class="figure">
          <img src="data:image/png;base64, %s" class="figure-img img-thumbnail input-image">
        </figure>        
        </html>
    """ %(o_img_data))
# <figure class="figure">
#           <img src="data:image/png;base64, %s" class="figure-img img-thumbnail input-image">
#         </figure>

@app.route("/")
def form(request):
    return HTMLResponse(
        """
        <h1>Neural Style Transfer!</h1>
        <p>Content and style images, will return content image styled like style.</p><br>
        <form action="/upload" method="post" enctype="multipart/form-data">
            <u>Select content and style image to upload:</u><br><p>
            Content <input type="file" name="content"><br><p>
            Style <input type="file" name="style"><br><p>
            Transfer! <input type="submit" value="Upload and transfer style">
        </form>
    """)




@app.route("/test")
def test(request):
    return HTMLResponse(
        """
        <h1>A test page?</h1>
        """)

@app.route("/form")
def redirect_to_homepage(request):
    return RedirectResponse("/")


if __name__ == "__main__":
    if "serve" in sys.argv:
        port = int(os.environ.get("PORT", 8008)) 
        uvicorn.run(app, host="0.0.0.0", port=port)

# -



StarcoderdataJupyterScriptsDedupFiltered