use kornia io

app.py

@@ -1,30 +1,28 @@
 import gradio as gr
 import torch
 import kornia as K
+from kornia.core import Tensor
 from kornia.geometry.transform import resize
-import cv2
-import numpy as np
-from torchvision import transforms
+
 from torchvision.utils import make_grid
 
+eps: float = 0.01
 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 read_image(f_name: str) -> Tensor:
+  # load the image using the rust backend          
+  img: Tensor = K.io.load_image(file.name, K.io.ImageLoadType.RGB32)
+  img = img[None]  # 1xCxHxW / fp32 / [0, 1]
+  return resize(img,(50, 50))
 
-def predict(images, eps):
-  eps = float(eps)
-  f_names = [img.name for img in images]
-  images = [read_image(f) for f in f_names]
+def predict(images):
+  images = [read_image(f.name) 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])
+  grid_zca = make_grid(zca_images, nrow=3, normalize=True)
+  return K.tensor_to_image(grid_zca)
 
 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:
@@ -49,7 +47,7 @@ iface = gr.Interface(fn=predict,
                   'Carnation.jpg', 
                   'Orchid.jpg', 
                   'Peony.jpg'
-                  ], 0.01]]
+                  ]]]
              )
 
 if __name__ == "__main__":