Spaces:
Runtime error
Runtime error
File size: 3,206 Bytes
93b3b80 |
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 |
# **************************************************************************** #
# #
# ::: :::::::: #
# style_model.py :+: :+: :+: #
# +:+ +:+ +:+ #
# By: ebennace <ebennace@student.42lausanne.c +#+ +:+ +#+ #
# +#+#+#+#+#+ +#+ #
# Created: 2022/11/15 13:08:56 by ebennace #+# #+# #
# Updated: 2022/11/15 13:08:57 by ebennace ### ########.fr #
# #
# **************************************************************************** #
from typing import Tuple
import numpy as np
from time import time
from tqdm.auto import tqdm
from tensorflow.keras.optimizers import Adam
from .vgg import create_list_of_vgg_layer
from .vgg import create_multi_output_model
from .init import init_style_target
from .init import init_generated_img
from .style_function import update_style
# ======================================================================= #
class StyleRecreationModel:
"""
A class for generating style representation of Input Image using the VGG19 model.
"""
# ======================================================================= #
def __init__(self):
"""
Initialize the StyleRecreationModel with a pre-trained VGG19 model and Adam optimizer.
"""
self.optimizer = Adam(learning_rate=0.02)
self.style_layers = create_list_of_vgg_layer()
self.num_style_layers = len(self.style_layers)
self.model = create_multi_output_model(self.style_layers)
# ======================================================================= #
def recreate_style(
self,
style_img_array : np.array,
num_epochs : int,
)->Tuple[np.array, float]:
"""
Generate a new image based on the style of the given input image.
Args:
`style_img_array` (np.array): The input style image as a NumPy array.
`num_epochs` (int): The number of epochs for the style transfer process.
Returns:
Tuple[np.array, float]: A tuple containing the generated image as a NumPy array and the total time taken.
"""
target_style = init_style_target(self.model, style_img_array)
self.generated_img = init_generated_img(style_img_array)
start = time()
for _ in tqdm(range(num_epochs)) :
update_style(
model=self.model,
style_target=target_style,
generated_img=self.generated_img,
optimizer=self.optimizer
)
end = time()
total_time = round(end-start, 2)
return (self.generated_img, total_time)
# ======================================================================= # |