import gradio as gr
from image_dataset import ImageDataset
from image_wgan import ImageWgan
import os
from os.path import exists
from PIL import Image
def init():
    generated_samples_folder = "."
    discriminator_saved_model = "discriminator64.model"
    generator_saved_model = "generator64.model"
    latent_space = 100
    image_wgan = ImageWgan(
        image_shape = (4,64,64),
        latent_space_dimension=latent_space,
        generator_saved_model=generator_saved_model if exists(generator_saved_model) else None,
        discriminator_saved_model=discriminator_saved_model if exists(discriminator_saved_model) else None
    )
    image_wgan.generate(
        sample_folder=generated_samples_folder
    )
    crop()

def crop():

    import generator
    res = 64
    if res != 0:
        results = "generated.png"
        img = Image.open(results)

        width,height = img.size


        top = 2
        bottom = 2
        for i in range(4):
            left = (res+2)*i +2
            right = width-(res+2)*i
            imgcrop = img.crop((left,top,left+res,res+2))


            imgcrop.save(str(i)+".png")     
           
init()






def gen():
    init()
    crop()
    img = Image.open("0.png")

    return img

iface = gr.Interface(
    fn=gen, 
    inputs=None, 
    outputs="image"),
    theme="darkhuggingface"
iface.launch(debug = True)