0xrushi commited on
Commit
658657b
1 Parent(s): 9d9224d

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +168 -0
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)