Spaces:
Running
on
CPU Upgrade
Running
on
CPU Upgrade
Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,168 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import gradio as gr
|
2 |
+
import numpy as np
|
3 |
+
import tensorflow as tf
|
4 |
+
from tensorflow import keras
|
5 |
+
from huggingface_hub import from_pretrained_keras
|
6 |
+
|
7 |
+
result_prefix = "paris_generated"
|
8 |
+
|
9 |
+
# Weights of the different loss components
|
10 |
+
total_variation_weight = 1e-6
|
11 |
+
style_weight = 1e-6
|
12 |
+
content_weight = 2.5e-8
|
13 |
+
|
14 |
+
# Dimensions of the generated picture.
|
15 |
+
width, height = keras.preprocessing.image.load_img(base_image_path).size
|
16 |
+
img_nrows = 400
|
17 |
+
img_ncols = int(width * img_nrows / height)
|
18 |
+
|
19 |
+
def preprocess_image(image_path):
|
20 |
+
# Util function to open, resize and format pictures into appropriate tensors
|
21 |
+
img = keras.preprocessing.image.load_img(
|
22 |
+
image_path, target_size=(img_nrows, img_ncols)
|
23 |
+
)
|
24 |
+
img = keras.preprocessing.image.img_to_array(img)
|
25 |
+
img = np.expand_dims(img, axis=0)
|
26 |
+
img = vgg19.preprocess_input(img)
|
27 |
+
return tf.convert_to_tensor(img)
|
28 |
+
|
29 |
+
def deprocess_image(x):
|
30 |
+
# Util function to convert a tensor into a valid image
|
31 |
+
x = x.reshape((img_nrows, img_ncols, 3))
|
32 |
+
# Remove zero-center by mean pixel
|
33 |
+
x[:, :, 0] += 103.939
|
34 |
+
x[:, :, 1] += 116.779
|
35 |
+
x[:, :, 2] += 123.68
|
36 |
+
# 'BGR'->'RGB'
|
37 |
+
x = x[:, :, ::-1]
|
38 |
+
x = np.clip(x, 0, 255).astype("uint8")
|
39 |
+
return x
|
40 |
+
|
41 |
+
# The gram matrix of an image tensor (feature-wise outer product)
|
42 |
+
|
43 |
+
def gram_matrix(x):
|
44 |
+
x = tf.transpose(x, (2, 0, 1))
|
45 |
+
features = tf.reshape(x, (tf.shape(x)[0], -1))
|
46 |
+
gram = tf.matmul(features, tf.transpose(features))
|
47 |
+
return gram
|
48 |
+
|
49 |
+
# The "style loss" is designed to maintain
|
50 |
+
# the style of the reference image in the generated image.
|
51 |
+
# It is based on the gram matrices (which capture style) of
|
52 |
+
# feature maps from the style reference image
|
53 |
+
# and from the generated image
|
54 |
+
|
55 |
+
def style_loss(style, combination):
|
56 |
+
S = gram_matrix(style)
|
57 |
+
C = gram_matrix(combination)
|
58 |
+
channels = 3
|
59 |
+
size = img_nrows * img_ncols
|
60 |
+
return tf.reduce_sum(tf.square(S - C)) / (4.0 * (channels ** 2) * (size ** 2))
|
61 |
+
|
62 |
+
# An auxiliary loss function
|
63 |
+
# designed to maintain the "content" of the
|
64 |
+
# base image in the generated image
|
65 |
+
|
66 |
+
def content_loss(base, combination):
|
67 |
+
return tf.reduce_sum(tf.square(combination - base))
|
68 |
+
|
69 |
+
# The 3rd loss function, total variation loss,
|
70 |
+
# designed to keep the generated image locally coherent
|
71 |
+
|
72 |
+
def total_variation_loss(x):
|
73 |
+
a = tf.square(
|
74 |
+
x[:, : img_nrows - 1, : img_ncols - 1, :] - x[:, 1:, : img_ncols - 1, :]
|
75 |
+
)
|
76 |
+
b = tf.square(
|
77 |
+
x[:, : img_nrows - 1, : img_ncols - 1, :] - x[:, : img_nrows - 1, 1:, :]
|
78 |
+
)
|
79 |
+
return tf.reduce_sum(tf.pow(a + b, 1.25))
|
80 |
+
|
81 |
+
def compute_loss(combination_image, base_image, style_reference_image):
|
82 |
+
input_tensor = tf.concat(
|
83 |
+
[base_image, style_reference_image, combination_image], axis=0
|
84 |
+
)
|
85 |
+
features = feature_extractor(input_tensor)
|
86 |
+
|
87 |
+
# Initialize the loss
|
88 |
+
loss = tf.zeros(shape=())
|
89 |
+
|
90 |
+
# Add content loss
|
91 |
+
layer_features = features[content_layer_name]
|
92 |
+
base_image_features = layer_features[0, :, :, :]
|
93 |
+
combination_features = layer_features[2, :, :, :]
|
94 |
+
loss = loss + content_weight * content_loss(
|
95 |
+
base_image_features, combination_features
|
96 |
+
)
|
97 |
+
# Add style loss
|
98 |
+
for layer_name in style_layer_names:
|
99 |
+
layer_features = features[layer_name]
|
100 |
+
style_reference_features = layer_features[1, :, :, :]
|
101 |
+
combination_features = layer_features[2, :, :, :]
|
102 |
+
sl = style_loss(style_reference_features, combination_features)
|
103 |
+
loss += (style_weight / len(style_layer_names)) * sl
|
104 |
+
|
105 |
+
# Add total variation loss
|
106 |
+
loss += total_variation_weight * total_variation_loss(combination_image)
|
107 |
+
return loss
|
108 |
+
|
109 |
+
# Build a VGG19 model loaded with pre-trained ImageNet weights
|
110 |
+
# model = vgg19.VGG19(weights="imagenet", include_top=False)
|
111 |
+
model = from_pretrained_keras("rushic24/keras-VGG19")
|
112 |
+
|
113 |
+
# Get the symbolic outputs of each "key" layer (we gave them unique names).
|
114 |
+
outputs_dict = dict([(layer.name, layer.output) for layer in model.layers])
|
115 |
+
|
116 |
+
# Set up a model that returns the activation values for every layer in
|
117 |
+
# VGG19 (as a dict).
|
118 |
+
feature_extractor = keras.Model(inputs=model.inputs, outputs=outputs_dict)
|
119 |
+
|
120 |
+
# List of layers to use for the style loss.
|
121 |
+
style_layer_names = [
|
122 |
+
"block1_conv1",
|
123 |
+
"block2_conv1",
|
124 |
+
"block3_conv1",
|
125 |
+
"block4_conv1",
|
126 |
+
"block5_conv1",
|
127 |
+
]
|
128 |
+
# The layer to use for the content loss.
|
129 |
+
content_layer_name = "block5_conv2"
|
130 |
+
|
131 |
+
@tf.function
|
132 |
+
def compute_loss_and_grads(combination_image, base_image, style_reference_image):
|
133 |
+
with tf.GradientTape() as tape:
|
134 |
+
loss = compute_loss(combination_image, base_image, style_reference_image)
|
135 |
+
grads = tape.gradient(loss, combination_image)
|
136 |
+
return loss, grads
|
137 |
+
|
138 |
+
optimizer = keras.optimizers.SGD(
|
139 |
+
keras.optimizers.schedules.ExponentialDecay(
|
140 |
+
initial_learning_rate=100.0, decay_steps=100, decay_rate=0.96
|
141 |
+
)
|
142 |
+
)
|
143 |
+
|
144 |
+
def get_imgs(base_image_path, style_reference_image_path):
|
145 |
+
base_image = preprocess_image(base_image_path)
|
146 |
+
style_reference_image = preprocess_image(style_reference_image_path)
|
147 |
+
combination_image = tf.Variable(preprocess_image(base_image_path))
|
148 |
+
|
149 |
+
iterations = 400
|
150 |
+
for i in range(1, iterations + 1):
|
151 |
+
loss, grads = compute_loss_and_grads(combination_image, base_image, style_reference_image)
|
152 |
+
optimizer.apply_gradients([(grads, combination_image)])
|
153 |
+
if i % 100 == 0:
|
154 |
+
print("Iteration %d: loss=%.2f" % (i, loss))
|
155 |
+
img = deprocess_image(combination_image.numpy())
|
156 |
+
|
157 |
+
return img
|
158 |
+
|
159 |
+
|
160 |
+
title = "Neural style transfer"
|
161 |
+
description = "Gradio Demo for Neural style transfer. To use it, simply upload a base image and a style image"
|
162 |
+
|
163 |
+
content = gr.inputs.Image(shape=None, image_mode="RGB", invert_colors=False, source="upload", tool="editor", type="filepath", label=None, optional=False)
|
164 |
+
style = gr.inputs.Image(shape=None, image_mode="RGB", invert_colors=False, source="upload", tool="editor", type="filepath", label=None, optional=False)
|
165 |
+
gr.Interface(get_imgs, inputs=[content, style], outputs=["image"],
|
166 |
+
title=title,
|
167 |
+
description=description,
|
168 |
+
examples=[["base.jpg", "style.jpg"]]).launch(enable_queue=True)
|