import gradio as gr
import QuckDrawGAN as qd
import numpy as np
from PIL import Image, ImageFilter, ImageEnhance

# Загрузим модель
generator_file = r'pretrained_output/models/generator.pt'
discriminator_file = r'pretrained_output/models/discriminator_fine_tuned.pt'

# Создаем объект модели
model = qd.Model(generator_file, discriminator_file)

# Функция для генерации изображения с учетом сида, нормализацией и изменением размера
def generate_image(n_images=16, seed=""):
    # Если сид не задан, не передаем его в модель
    if seed == "":
        seed = None
    best_image = model.generate(n_images, seed)  # Генерация с учетом сида

    # Нормализация: находим минимум и максимум в изображении
    best_image_min = np.min(best_image)
    best_image_max = np.max(best_image)

    # Нормализуем изображение, чтобы значения были в диапазоне от 0 до 255
    normalized_image = 255 * (best_image - best_image_min) / (best_image_max - best_image_min)

    # Преобразуем изображение в формат, подходящий для отображения
    pil_image = Image.fromarray(normalized_image.astype(np.uint8))  # Преобразуем в uint8 для отображения
    pil_image = pil_image.resize((256, 256), Image.Resampling.LANCZOS)  # Ресайз изображения до 256x256

    # Применяем размытие
    pil_image = pil_image.filter(ImageFilter.GaussianBlur(radius=1))  # Фильтр размытия с радиусом 2

    # Применяем контрастность
    enhancer = ImageEnhance.Contrast(pil_image)  # Создаем объект для изменения контрастности
    pil_image = enhancer.enhance(3)  # Увеличиваем контрастность на 1.5

    return pil_image

# Создаем интерфейс Gradio с помощью Blocks (для большей гибкости)
with gr.Blocks() as interface:
    gr.Markdown("# Генератор рисунков с использованием предобученной модели QuckDrawGAN (Duck)")

    with gr.Row():
        # Блок для изображения сверху
        generated_image = gr.Image(type="pil", label="Сгенерированное изображение", elem_id="generated_image", scale=2)  # Увеличиваем масштаб изображения

    with gr.Row():
        # Блок параметров и кнопки
        with gr.Column():
            seed_input = gr.Textbox(value="", label="Сид (опционально)", interactive=True)
            num_images = gr.Slider(minimum=1, maximum=1024, value=32, label="Количество изображений для генерации, из них модель 'критик' выберет самое лучшее", interactive=True, step=1)
        
        # Кнопка генерации изображения справа
        generate_button = gr.Button("Сгенерировать")

    # Логика для автогенерации при изменении параметров
    seed_input.change(generate_image, inputs=[num_images, seed_input], outputs=generated_image)
    num_images.change(generate_image, inputs=[num_images, seed_input], outputs=generated_image)

    # Логика для кнопки генерации
    generate_button.click(generate_image, inputs=[num_images, seed_input], outputs=generated_image)

    # Автогенерация при старте
    interface.load(generate_image, inputs=[num_images, seed_input], outputs=generated_image)

    # Стилизация блока изображения (увеличение размера блока)
    interface.css = """
    #generated_image {
        width: 400px;
        height: 400px;
        margin-top: 20px;
    }
    """

# Запуск интерфейса
interface.launch()