import gradio as gr
import torch
import kornia as K
from kornia.geometry.transform import resize
import numpy as np
from torchvision import transforms
from torchvision.utils import make_grid
from PIL import Image

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def read_image(img):
    image_to_tensor = transforms.ToTensor()
    if isinstance(img, np.ndarray):
        img = Image.fromarray(img)
    elif isinstance(img, str):
        img = Image.open(img).convert('RGB')
    img_tensor = image_to_tensor(img)
    resized_image = resize(img_tensor.unsqueeze(0), (50, 50)).squeeze(0)
    return resized_image

def predict(images, eps):
    eps = float(eps)
    images = [read_image(img) for img in images]
    images = torch.stack(images, dim=0).to(device)
    zca = K.enhance.ZCAWhitening(eps=eps, compute_inv=True)
    zca.fit(images)
    zca_images = zca(images)
    grid_zca = make_grid(zca_images, nrow=3, normalize=True).cpu().numpy()
    return np.transpose(grid_zca, [1, 2, 0])

def load_example_images():
    return example_images

title = 'ZCA Whitening with Kornia!'
description = '''[ZCA Whitening](https://paperswithcode.com/method/zca-whitening) is an image preprocessing method that leads to a transformation of data such that the covariance matrix is the identity matrix, leading to decorrelated features:
*Note that you can upload only image files, e.g. jpg, png etc and there should be at least 2 images!* 
Learn more about [ZCA Whitening and Kornia](https://kornia.readthedocs.io/en/latest/enhance.zca.html)'''

example_images = ['irises.jpg', 'roses.jpg', 'sunflower.jpg', 'violets.jpg', 'chamomile.jpg', 
                  'tulips.jpg', 'Alstroemeria.jpg', 'Carnation.jpg', 'Orchid.jpg', 'Peony.jpg']

with gr.Blocks(title=title) as demo:
    gr.Markdown(f"# {title}")
    gr.Markdown(description)
    
    with gr.Row():
        input_images = gr.Files(label="Input Images")
        eps_slider = gr.Slider(minimum=0.01, maximum=1, value=0.01, label="Epsilon")
    
    output_image = gr.Image(label="ZCA Whitened Images")
    
    submit_button = gr.Button("Apply ZCA Whitening")
    submit_button.click(fn=predict, inputs=[input_images, eps_slider], outputs=output_image)
    
    gr.Markdown("## Example Images")
    example_gallery = gr.Gallery(value=example_images, label="Example Images", columns=5, height="auto")
    load_examples_button = gr.Button("Load Example Images")
    load_examples_button.click(fn=load_example_images, inputs=[], outputs=[input_images])

if __name__ == "__main__":
    demo.launch(show_error=True)