Spaces:
Runtime error
Runtime error
| # **************************************************************************** # | |
| # # | |
| # ::: :::::::: # | |
| # 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) | |
| # ======================================================================= # |