Update kornia_aug.py

kornia_aug.py

@@ -18,21 +18,27 @@ def set_transform(content):
         transform = nn.Sequential()
     return transform
 
+st.set_page_config(page_title="Kornia Augmentations Demo", layout="wide")
+
 st.markdown("# Kornia Augmentations Demo")
 st.sidebar.markdown(
     "[Kornia](https://github.com/kornia/kornia) is a *differentiable* computer vision library for PyTorch."
 )
-uploaded_file = st.sidebar.file_uploader("Choose a file")
+
+uploaded_file = st.sidebar.file_uploader("Choose a file", type=['png', 'jpg', 'jpeg'])
 if uploaded_file is not None:
     im = Image.open(uploaded_file)
 else:
     im = Image.open("./images/pretty_bird.jpg")
+
 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)
+
 gpu = st.sidebar.checkbox("Use GPU!", value=True)
 if not gpu:
     st.sidebar.markdown("With Kornia you do ops on the GPU!")
@@ -93,6 +99,7 @@ content = st_ace(
     auto_update=False,
     readonly=readonly,
 )
+
 if content:
     transform = set_transform(content)
 
@@ -100,8 +107,8 @@ process = st.button("Next Batch")
 
 # Fake dataloader
 image_batch = torch.stack(batch_size * [image])
-
 image_batch = image_batch.to(device)
+
 transformeds = None
 try:
     transformeds = transform(image_batch)
@@ -109,7 +116,6 @@ except Exception as e:
     st.write(f"There was an error: {e}")
 
 cols = st.columns(4)
-
 if transformeds is not None:
     for i, x in enumerate(transformeds):
         i = i % 4