Update kornia_aug.py

kornia_aug.py

@@ -1,17 +1,16 @@
 import streamlit as st
 import kornia
-from torch import nn
 import torch
+from torch import nn
 from torchvision.transforms import functional as F
 from torchvision.utils import make_grid
 from streamlit_ace import st_ace
 from PIL import Image
 
-IS_LOCAL = False #Change this
+IS_LOCAL = False  # Change this
 
-@st.cache(suppress_st_warning=True)
+@st.cache_data
 def set_transform(content):
-    #     st.write("set transform")
     try:
         transform = eval(content, {"kornia": kornia, "nn": nn}, None)
     except Exception as e:
@@ -32,9 +31,8 @@ scaler = int(im.height / 2)
 st.sidebar.image(im, caption="Input Image", width=256)
 image = F.pil_to_tensor(im).float() / 255
 
-
 # batch size is just for show
-batch_size = st.sidebar.slider("batch_size", min_value=4, max_value=16,value=8)
+batch_size = st.sidebar.slider("batch_size", min_value=4, max_value=16, value=8)
 gpu = st.sidebar.checkbox("Use GPU!", value=True)
 if not gpu:
     st.sidebar.markdown("With Kornia you do ops on the GPU!")
@@ -48,7 +46,7 @@ else:
         device = torch.device("cpu")
     else:
         st.sidebar.markdown("Running on GPU~")
-        device = torch.device("cuda:0")
+        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
 predefined_transforms = [
     """
@@ -75,7 +73,7 @@ nn.Sequential(
     kornia.augmentation.RandomHorizontalFlip(p=0.7),
     kornia.augmentation.RandomGrayscale(p=0.5),
 )
-""",
+"""
 ]
 
 selected_transform = st.selectbox(
@@ -96,36 +94,22 @@ content = st_ace(
     readonly=readonly,
 )
 if content:
-    #     st.write(content)
     transform = set_transform(content)
 
-# st.write(transform)
-
-# with st.echo():
-#     transform = nn.Sequential(
-#        K.RandomAffine(360),
-#        K.ColorJitter(0.2, 0.3, 0.2, 0.3)
-#     )
-
 process = st.button("Next Batch")
 
 # Fake dataloader
 image_batch = torch.stack(batch_size * [image])
 
-
-image_batch.to(device)
+image_batch = image_batch.to(device)
 transformeds = None
 try:
     transformeds = transform(image_batch)
 except Exception as e:
     st.write(f"There was an error: {e}")
-    
-
-
 
 cols = st.columns(4)
 
-# st.image(F.to_pil_image(make_grid(transformeds)))
 if transformeds is not None:
     for i, x in enumerate(transformeds):
         i = i % 4
@@ -136,7 +120,4 @@ st.markdown(
 )
 st.markdown(
     "Kornia can do a lot more than augmentations~ [Check it out](https://kornia.readthedocs.io/en/latest/get-started/introduction.html)"
-)
-# if process:
-#     pass
-
+)