import gradio as gr
from image_dataset import ImageDataset
from image_wgan import ImageWgan
import os
from os.path import exists
def init():
    generated_samples_folder = "results64"
    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():
    from PIL import Image
    import generator
    res = 64
    if res != 0:
        results = "results64/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("results64/"+str(i)+".png")        
init()






def generate(name):
    init()
    crop()
    img = Image.open("results64/0.png")
    return img

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