File size: 4,702 Bytes
760312f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 |
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
from tensorflow.keras.preprocessing import image as tf_image
from tensorflow.keras.applications import VGG19
from tensorflow.keras.models import Model
from tensorflow.keras import layers
from tensorflow.keras.optimizers import Adam
# Load a pre-trained VGG19 model without the fully connected layers (used for feature extraction)
base_model = VGG19(weights="imagenet", include_top=False)
# Specify the layers to use for style and content representations
style_layers = ["block1_conv1", "block2_conv1", "block3_conv1", "block4_conv1", "block5_conv1"]
content_layer = "block4_conv2"
# Create a model that extracts style and content features
style_extractor = Model(inputs=base_model.input, outputs=[base_model.get_layer(layer).output for layer in style_layers])
content_extractor = Model(inputs=base_model.input, outputs=base_model.get_layer(content_layer).output)
# Define a function to compute the Gram matrix for style representation
def gram_matrix(input_tensor):
result = tf.linalg.einsum("bijc,bijd->bcd", input_tensor, input_tensor)
input_shape = tf.shape(input_tensor)
num_locations = tf.cast(input_shape[1]*input_shape[2], tf.float32)
return result / num_locations
# Define a custom loss function that computes the style loss
def style_loss(style_targets, predicted_styles):
loss = 0
for style_target, predicted_style in zip(style_targets, predicted_styles):
loss += tf.reduce_mean(tf.square(gram_matrix(style_target) - gram_matrix(predicted_style)))
return loss
# Define a custom loss function that computes the content loss
def content_loss(content_target, predicted_content):
return tf.reduce_mean(tf.square(content_target - predicted_content))
# Load your input and style images
input_image_path = "input_image.jpg"
style_image_path = "style_image.jpg"
input_image = tf_image.load_img(input_image_path)
style_image = tf_image.load_img(style_image_path)
input_image = tf_image.img_to_array(input_image)
style_image = tf_image.img_to_array(style_image)
# Preprocess the images (VGG19 requires specific preprocessing)
input_image = tf_image.smart_resize(input_image, (256, 256))
style_image = tf_image.smart_resize(style_image, (256, 256))
input_image = tf_image.img_to_array(input_image)
style_image = tf_image.img_to_array(style_image)
input_image = tf.keras.applications.vgg19.preprocess_input(input_image)
style_image = tf.keras.applications.vgg19.preprocess_input(style_image)
input_image = np.expand_dims(input_image, axis=0)
style_image = np.expand_dims(style_image, axis=0)
# Define a variable to store the generated image and create a TensorFlow variable for it
generated_image = tf.Variable(input_image, dtype=tf.float32)
# Define optimizer and hyperparameters
optimizer = Adam(learning_rate=10.0)
# Number of iterations for optimization
num_iterations = 1000
# Extract style and content features from the style and input images
style_features = style_extractor(style_image)
content_features = content_extractor(input_image)
# Define target style features (the same style for all layers)
style_targets = [style_extractor(tf.constant(style_image)) for _ in style_layers]
# Main optimization loop
for iteration in range(num_iterations):
with tf.GradientTape() as tape:
# Extract features from the generated image
generated_features = style_extractor(generated_image)
# Compute style loss and content loss
current_style_loss = style_loss(style_targets, generated_features)
current_content_loss = content_loss(content_features, generated_features[-1])
# Total loss as a combination of style and content loss
total_loss = current_style_loss + current_content_loss
# Compute gradients
gradients = tape.gradient(total_loss, generated_image)
# Update the generated image using the gradients
optimizer.apply_gradients([(gradients, generated_image)])
# Clip pixel values to the [0, 255] range
generated_image.assign(tf.clip_by_value(generated_image, clip_value_min=0.0, clip_value_max=255.0))
# Print the progress
if iteration % 100 == 0:
print(f"Iteration {iteration}, Total loss: {total_loss}")
# Convert the final generated image to a NumPy array
final_image = tf_image.img_to_array(generated_image[0])
# Clip pixel values to the [0, 255] range and cast to uint8
final_image = np.clip(final_image, 0, 255).astype(np.uint8)
# Save the final image
final_image_path = "enhanced_image.jpg"
tf.keras.preprocessing.image.save_img(final_image_path, final_image[0])
# Display the final enhanced image
plt.imshow(final_image[0])
plt.axis("off")
plt.show()
|