ytkoa/QuckDrawGAN

1import gradio as gr2import QuckDrawGAN as qd3import numpy as np4from PIL import Image, ImageFilter, ImageEnhance5 6# Загрузим модель7generator_file = r'pretrained_output/models/generator.pt'8discriminator_file = r'pretrained_output/models/discriminator_fine_tuned.pt'9 10# Создаем объект модели11model = qd.Model(generator_file, discriminator_file)12 13# Функция для генерации изображения с учетом сида, нормализацией и изменением размера14def generate_image(n_images=16, seed=""):15    # Если сид не задан, не передаем его в модель16    if seed == "":17        seed = None18    best_image = model.generate(n_images, seed)  # Генерация с учетом сида19 20    # Нормализация: находим минимум и максимум в изображении21    best_image_min = np.min(best_image)22    best_image_max = np.max(best_image)23 24    # Нормализуем изображение, чтобы значения были в диапазоне от 0 до 25525    normalized_image = 255 * (best_image - best_image_min) / (best_image_max - best_image_min)26 27    # Преобразуем изображение в формат, подходящий для отображения28    pil_image = Image.fromarray(normalized_image.astype(np.uint8))  # Преобразуем в uint8 для отображения29    pil_image = pil_image.resize((256, 256), Image.Resampling.LANCZOS)  # Ресайз изображения до 256x25630 31    # Применяем размытие32    pil_image = pil_image.filter(ImageFilter.GaussianBlur(radius=1))  # Фильтр размытия с радиусом 233 34    # Применяем контрастность35    enhancer = ImageEnhance.Contrast(pil_image)  # Создаем объект для изменения контрастности36    pil_image = enhancer.enhance(3)  # Увеличиваем контрастность на 1.537 38    return pil_image39 40# Создаем интерфейс Gradio с помощью Blocks (для большей гибкости)41with gr.Blocks() as interface:42    gr.Markdown("# Генератор рисунков с использованием предобученной модели QuckDrawGAN (Duck)")43 44    with gr.Row():45        # Блок для изображения сверху46        generated_image = gr.Image(type="pil", label="Сгенерированное изображение", elem_id="generated_image", scale=2)  # Увеличиваем масштаб изображения47 48    with gr.Row():49        # Блок параметров и кнопки50        with gr.Column():51            seed_input = gr.Textbox(value="", label="Сид (опционально)", interactive=True)52            num_images = gr.Slider(minimum=1, maximum=1024, value=32, label="Количество изображений для генерации, из них модель 'критик' выберет самое лучшее", interactive=True, step=1)53        54        # Кнопка генерации изображения справа55        generate_button = gr.Button("Сгенерировать")56 57    # Логика для автогенерации при изменении параметров58    seed_input.change(generate_image, inputs=[num_images, seed_input], outputs=generated_image)59    num_images.change(generate_image, inputs=[num_images, seed_input], outputs=generated_image)60 61    # Логика для кнопки генерации62    generate_button.click(generate_image, inputs=[num_images, seed_input], outputs=generated_image)63 64    # Автогенерация при старте65    interface.load(generate_image, inputs=[num_images, seed_input], outputs=generated_image)66 67    # Стилизация блока изображения (увеличение размера блока)68    interface.css = """69    #generated_image {70        width: 400px;71        height: 400px;72        margin-top: 20px;73    }74    """75 76# Запуск интерфейса77interface.launch()78