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

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

def read_image(f_name):
      image_to_tensor = transforms.ToTensor()
      img = image_to_tensor(cv2.imread(f_name, cv2.IMREAD_COLOR))
      resized_image = resize(img,(50, 50))
      return resized_image

def predict(images, eps):
  eps = float(eps)
  f_names = [img.name for img in images]
  images = [read_image(f) for f in f_names]
  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])

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 sjould be atleast 2 images!* 
Learn more about [ZCA Whitening and Kornia](https://kornia.readthedocs.io/en/latest/_modules/kornia/enhance/zca.html)'''

iface = gr.Interface(fn=predict, 
             inputs=['files', gr.Slider(0.01, 1)],
             outputs=gr.Image(),
             allow_flagging="never",
             title=title,
             description=description,
             examples=[[
                 [
                  'irises.jpg', 
                  'roses.jpg', 
                  'sunflower.jpg',
                  'violets.jpg', 
                  'chamomile.jpg', 
                  'tulips.jpg', 
                  'Alstroemeria.jpg', 
                  'Carnation.jpg', 
                  'Orchid.jpg', 
                  'Peony.jpg'
                  ], 0.01]]
             )

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