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, ImageEnhance
import numpy
print(gr.__version__)
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, top + res))  # Changed res+2 to res

            # Resize the image using nearest neighbor interpolation
            imgcrop = imgcrop.resize((res, res), Image.NEAREST)

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





import numpy

def gen(seed):
    numpy.random.seed(int(seed))
    init()
    crop()
    imgArr = []
    for i in range(4):
        img = Image.open(str(i)+".png")
        imgArr.append(img)
    return imgArr

iface = gr.Interface(
    fn=gen, 
    inputs=gr.Slider(0,1000,500,step = 0.01),
    outputs=gr.Gallery(label="Generated Skins"),
    title = "Minecraft Skin Generator",
    css = "html{image-rendering:pixelated}",
    debug = True,
)
iface.launch(width=64,height=64,enable_queue=True)