add initial code for app
This commit is contained in:
s486797
parent
c54f43c735
commit
074482998d
1
neural_style_app/.gitignore
vendored
Normal file
1
neural_style_app/.gitignore
vendored
Normal file
@ -0,0 +1 @@
|
||||
__pycache__/
|
||||
99
neural_style_app/app.py
Normal file
99
neural_style_app/app.py
Normal file
@ -0,0 +1,99 @@
|
||||
from flask import Flask, render_template, request, redirect, url_for, send_file, jsonify
|
||||
from mode_style_transfer import StyleTransferModel, save_image, StyleTransferVisualizer
|
||||
from PIL import Image
|
||||
import io
|
||||
import torch
|
||||
from torchvision.models import vgg19, VGG19_Weights
|
||||
import torchvision.transforms as transforms
|
||||
import os
|
||||
import matplotlib.pyplot as plt
|
||||
import base64
|
||||
|
||||
app = Flask(__name__)
|
||||
|
||||
# Image transformation
|
||||
imsize = 512 if torch.cuda.is_available() else 128
|
||||
loader = transforms.Compose([
|
||||
transforms.Resize(imsize),
|
||||
transforms.ToTensor()
|
||||
])
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
visualizations = []
|
||||
|
||||
def image_loader(image_bytes):
|
||||
image = Image.open(io.BytesIO(image_bytes))
|
||||
image = loader(image).unsqueeze(0)
|
||||
return image.to(device, torch.float)
|
||||
|
||||
def tensor_to_image(tensor):
|
||||
image = tensor.clone().detach().squeeze(0)
|
||||
image = transforms.ToPILImage()(image)
|
||||
return image
|
||||
|
||||
def image_to_base64(image):
|
||||
img_io = io.BytesIO()
|
||||
image.save(img_io, 'JPEG')
|
||||
img_io.seek(0)
|
||||
return base64.b64encode(img_io.getvalue()).decode('utf-8')
|
||||
|
||||
|
||||
@app.route('/', methods=['GET', 'POST'])
|
||||
def index():
|
||||
if request.method == 'POST':
|
||||
content_image_file = request.files['content_image']
|
||||
style_image_file = request.files['style_image']
|
||||
|
||||
# Load images directly from the uploaded files
|
||||
content_image = Image.open(content_image_file)
|
||||
style_image = Image.open(style_image_file)
|
||||
|
||||
# Pass the images to the StyleTransferModel
|
||||
style_transfer = StyleTransferModel(content_image, style_image)
|
||||
output = style_transfer.run_style_transfer()
|
||||
|
||||
# Convert the output tensor to an image
|
||||
output_image = tensor_to_image(output)
|
||||
|
||||
# Convert the image to Base64 for JSON response
|
||||
image_base64 = image_to_base64(output_image)
|
||||
|
||||
return jsonify({'image': image_base64})
|
||||
|
||||
return render_template('index.html')
|
||||
|
||||
@app.route('/visualize', methods=['POST'])
|
||||
def visualize():
|
||||
cnn = vgg19(weights=VGG19_Weights.DEFAULT).features.to(device).eval()
|
||||
|
||||
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)
|
||||
|
||||
content_image_bytes = visualizations[0] # The last saved content image
|
||||
content_image = image_loader(content_image_bytes)
|
||||
|
||||
style_transfer = StyleTransferModel(content_image, content_image)
|
||||
|
||||
# Running the model for visualization purpose
|
||||
input_img = content_image.clone().requires_grad_(True)
|
||||
|
||||
model, _, _ = style_transfer.get_style_model_and_losses(
|
||||
cnn, cnn_normalization_mean, cnn_normalization_std, content_image, content_image)
|
||||
|
||||
layer_visualizations = []
|
||||
|
||||
# Run the image through each layer and store the output
|
||||
for i, layer in enumerate(model):
|
||||
input_img = layer(input_img)
|
||||
with torch.no_grad():
|
||||
output_image = tensor_to_image(input_img.clamp(0, 1))
|
||||
img_io = io.BytesIO()
|
||||
output_image.save(img_io, 'JPEG')
|
||||
img_io.seek(0)
|
||||
layer_visualizations.append(img_io.getvalue()) # Save the image bytes
|
||||
|
||||
return render_template('visualize.html', visualizations=layer_visualizations)
|
||||
|
||||
if __name__ == '__main__':
|
||||
app.run(debug=True)
|
||||
258
neural_style_app/mode_style_transfer.py
Normal file
258
neural_style_app/mode_style_transfer.py
Normal file
@ -0,0 +1,258 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch.optim as optim
|
||||
|
||||
from PIL import Image
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
import torchvision.transforms as transforms
|
||||
from torchvision.models import vgg19, VGG19_Weights
|
||||
|
||||
from torchvision import models
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
torch.set_default_device(device)
|
||||
|
||||
def image_loader(image):
|
||||
#image = Image.open(image_name)
|
||||
# fake batch dimension required to fit network's input dimensions
|
||||
|
||||
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.ToTensor()]) # transform it into a torch tensor
|
||||
image = loader(image).unsqueeze(0)
|
||||
return image.to(device, torch.float)
|
||||
|
||||
def save_image(tensor, path):
|
||||
image = tensor.clone().detach()
|
||||
image = image.squeeze(0)
|
||||
image = transforms.ToPILImage()(image)
|
||||
image.save(path)
|
||||
|
||||
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) # resize 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
|
||||
|
||||
#cnn = vgg19(weights=VGG19_Weights.DEFAULT).features.eval()
|
||||
|
||||
#cnn_normalization_mean = torch.tensor([0.485, 0.456, 0.406])
|
||||
#cnn_normalization_std = torch.tensor([0.229, 0.224, 0.225])
|
||||
|
||||
# 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
|
||||
|
||||
# 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']
|
||||
|
||||
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):
|
||||
# normalization module
|
||||
normalization = Normalization(normalization_mean, normalization_std)
|
||||
|
||||
# just in order to have an iterable access to or list of content/style
|
||||
# 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])
|
||||
return optimizer
|
||||
|
||||
class StyleTransferModel:
|
||||
def __init__(self, content_img, style_img, num_steps=300, style_weight=1000000, content_weight=1):
|
||||
self.content_img = content_img
|
||||
self.style_img = style_img.resize(content_img.size)
|
||||
#self.style_img = self.style_img.resize(self.content_img.size)
|
||||
self.style_img = image_loader(self.style_img)
|
||||
self.content_img = image_loader(self.content_img)
|
||||
self.input_img = self.content_img.clone()
|
||||
self.num_steps = num_steps
|
||||
self.style_weight = style_weight
|
||||
self.content_weight = content_weight
|
||||
|
||||
self.cnn = vgg19(weights=VGG19_Weights.DEFAULT).features.to(device).eval()
|
||||
self.cnn_normalization_mean = torch.tensor([0.485, 0.456, 0.406]).to(device)
|
||||
self.cnn_normalization_std = torch.tensor([0.229, 0.224, 0.225]).to(device)
|
||||
|
||||
def run_style_transfer(self):
|
||||
print('Building the style transfer model..')
|
||||
model, style_losses, content_losses = get_style_model_and_losses(
|
||||
self.cnn, self.cnn_normalization_mean, self.cnn_normalization_std,
|
||||
self.style_img, self.content_img)
|
||||
|
||||
self.input_img.requires_grad_(True)
|
||||
model.eval()
|
||||
model.requires_grad_(False)
|
||||
optimizer = get_input_optimizer(self.input_img)
|
||||
|
||||
print('Optimizing..')
|
||||
run = [0]
|
||||
while run[0] <= self.num_steps:
|
||||
def closure():
|
||||
with torch.no_grad():
|
||||
self.input_img.clamp_(0, 1)
|
||||
|
||||
optimizer.zero_grad()
|
||||
model(self.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 *= self.style_weight
|
||||
content_score *= self.content_weight
|
||||
|
||||
loss = style_score + content_score
|
||||
loss.backward()
|
||||
|
||||
run[0] += 1
|
||||
if run[0] % 50 == 0:
|
||||
print(f"run {run[0]}:")
|
||||
print(f'Style Loss : {style_score.item():4f} Content Loss: {content_score.item():4f}')
|
||||
print()
|
||||
|
||||
return style_score + content_score
|
||||
|
||||
optimizer.step(closure)
|
||||
|
||||
with torch.no_grad():
|
||||
self.input_img.clamp_(0, 1)
|
||||
|
||||
return self.input_img
|
||||
|
||||
|
||||
class StyleTransferVisualizer(StyleTransferModel):
|
||||
def __init__(self, content_img, style_img):
|
||||
super().__init__(content_img, style_img)
|
||||
self.model_layers = self.get_model_layers()
|
||||
|
||||
def get_model_layers(self):
|
||||
cnn = models.vgg19(pretrained=True).features.to(self.device).eval()
|
||||
model_layers = []
|
||||
i = 0
|
||||
for layer in cnn.children():
|
||||
if isinstance(layer, torch.nn.Conv2d):
|
||||
i += 1
|
||||
model_layers.append((f'conv_{i}', layer))
|
||||
return model_layers
|
||||
|
||||
def visualize_layers(self):
|
||||
fig, axs = plt.subplots(len(self.model_layers), 3, figsize=(15, 20))
|
||||
|
||||
input_img = self.content_img.clone().detach()
|
||||
|
||||
for idx, (name, layer) in enumerate(self.model_layers):
|
||||
input_img = layer(input_img)
|
||||
axs[idx, 0].imshow(self.content_img.squeeze(0).permute(1, 2, 0).cpu().numpy())
|
||||
axs[idx, 0].set_title("Original Image")
|
||||
axs[idx, 0].axis('off')
|
||||
|
||||
axs[idx, 1].imshow(input_img.squeeze(0).permute(1, 2, 0).cpu().detach().numpy())
|
||||
axs[idx, 1].set_title(f"After {name}")
|
||||
axs[idx, 1].axis('off')
|
||||
|
||||
combined = input_img.clone()
|
||||
combined += self.style_img.squeeze(0)
|
||||
axs[idx, 2].imshow(combined.permute(1, 2, 0).cpu().detach().numpy())
|
||||
axs[idx, 2].set_title(f"Combined (Content + Style) after {name}")
|
||||
axs[idx, 2].axis('off')
|
||||
|
||||
plt.tight_layout()
|
||||
plt.show()
|
||||
71
neural_style_app/templates/index.html
Normal file
71
neural_style_app/templates/index.html
Normal file
@ -0,0 +1,71 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>Style Transfer</title>
|
||||
<style>
|
||||
/* Style the container to ensure it doesn't take up the whole page */
|
||||
#result-container {
|
||||
margin-top: 20px;
|
||||
border: 1px solid #ccc;
|
||||
padding: 10px;
|
||||
display: none; /* Hide the container initially */
|
||||
}
|
||||
|
||||
img {
|
||||
max-width: 100%;
|
||||
height: auto;
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<h1>Style Transfer</h1>
|
||||
<form id="upload-form" action="/" method="POST" enctype="multipart/form-data">
|
||||
<label for="content_image">Upload Content Image:</label>
|
||||
<input type="file" id="content_image" name="content_image" required><br><br>
|
||||
|
||||
<label for="style_image">Upload Style Image:</label>
|
||||
<input type="file" id="style_image" name="style_image" required><br><br>
|
||||
|
||||
<input type="submit" value="Submit">
|
||||
</form>
|
||||
|
||||
<!-- Container for the resulting image -->
|
||||
<div id="result-container">
|
||||
<h2>Resulting Image:</h2>
|
||||
<div id="image-container"></div>
|
||||
<a href="{{ url_for('visualize') }}">
|
||||
<button>Visualize Layers</button>
|
||||
</a>
|
||||
</div>
|
||||
|
||||
<script>
|
||||
// JavaScript to handle form submission and display the image
|
||||
document.getElementById('upload-form').addEventListener('submit', async function(event) {
|
||||
event.preventDefault(); // Prevent the form from submitting the traditional way
|
||||
|
||||
const formData = new FormData(this);
|
||||
|
||||
const response = await fetch('/', {
|
||||
method: 'POST',
|
||||
body: formData
|
||||
});
|
||||
|
||||
if (response.ok) {
|
||||
const data = await response.json();
|
||||
|
||||
const imageElement = document.createElement('img');
|
||||
imageElement.src = 'data:image/jpeg;base64,' + data.image;
|
||||
|
||||
const imageContainer = document.getElementById('image-container');
|
||||
imageContainer.innerHTML = ''; // Clear any previous image
|
||||
imageContainer.appendChild(imageElement);
|
||||
|
||||
const resultContainer = document.getElementById('result-container');
|
||||
resultContainer.style.display = 'block'; // Show the container
|
||||
}
|
||||
});
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
25
neural_style_app/templates/visualize.html
Normal file
25
neural_style_app/templates/visualize.html
Normal file
@ -0,0 +1,25 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>Visualize Layers</title>
|
||||
</head>
|
||||
<body>
|
||||
<h1>Layer Visualizations</h1>
|
||||
<p>Select a layer to view the image before and after processing through that layer:</p>
|
||||
|
||||
{% for i in range(visualizations|length) %}
|
||||
<div>
|
||||
<h2>Layer {{ i + 1 }}</h2>
|
||||
<button onclick="document.getElementById('img_before').src='{{ url_for(show_image, index=i) }}';">
|
||||
View Layer {{ i + 1 }}
|
||||
</button>
|
||||
</div>
|
||||
{% endfor %}
|
||||
|
||||
<h2>Layer Output:</h2>
|
||||
<img id="img_before" src="" alt="Layer Image Output" style="max-width: 100%; height: auto;">
|
||||
</body>
|
||||
</html>
|
||||
|
||||
@ -2,7 +2,7 @@
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 135,
|
||||
"execution_count": 1,
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
@ -69,7 +69,7 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 136,
|
||||
"execution_count": 2,
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
@ -91,11 +91,22 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 137,
|
||||
"execution_count": 3,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"True"
|
||||
]
|
||||
},
|
||||
"execution_count": 3,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"#torch.cuda.is_available()"
|
||||
"torch.cuda.is_available()"
|
||||
]
|
||||
},
|
||||
{
|
||||
@ -113,7 +124,7 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 138,
|
||||
"execution_count": 4,
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
@ -123,6 +134,26 @@
|
||||
"torch.set_default_device(device)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 5,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"device(type='cuda')"
|
||||
]
|
||||
},
|
||||
"execution_count": 5,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"device"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
@ -150,7 +181,7 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 154,
|
||||
"execution_count": 6,
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
@ -181,6 +212,26 @@
|
||||
" \"we need to import style and content images of the same size\"\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 7,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"512"
|
||||
]
|
||||
},
|
||||
"execution_count": 7,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"imsize"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
|
||||
Loading…
Reference in New Issue
Block a user