Spaces:
Running
on
CPU Upgrade
Running
on
CPU Upgrade
File size: 3,730 Bytes
d72ec76 9acfc48 d72ec76 9acfc48 d72ec76 9acfc48 d72ec76 9acfc48 d72ec76 9acfc48 d72ec76 9acfc48 d72ec76 9acfc48 d72ec76 f23b303 9acfc48 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 |
import streamlit as st
import kornia
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
@st.cache_data
def set_transform(content):
try:
transform = eval(content, {"kornia": kornia, "nn": nn}, None)
except Exception as e:
st.write(f"There was an error: {e}")
transform = nn.Sequential()
return transform
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")
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!")
device = torch.device("cpu")
else:
if not IS_LOCAL:
st.sidebar.markdown("(GPU Not available on hosted demo, try on your local!)")
# Credits
st.sidebar.caption("Demo made by [Ceyda Cinarel](https://linktr.ee/ceydai)")
st.sidebar.markdown("Clone [Code](https://github.com/cceyda/kornia-demo)")
device = torch.device("cpu")
else:
st.sidebar.markdown("Running on GPU~")
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
predefined_transforms = [
"""
nn.Sequential(
kornia.augmentation.RandomAffine(degrees=360,p=0.5),
kornia.augmentation.ColorJitter(brightness=0.2, contrast=0.3, saturation=0.2, hue=0.3, p=1)
)
# p=0.5 is the probability of applying the transformation
""",
"""
nn.Sequential(
kornia.augmentation.RandomErasing(scale=(.4, .8), ratio=(.3, 1/.3), p=0.5),
)
""",
"""
nn.Sequential(
kornia.augmentation.RandomErasing(scale=(.4, .8), ratio=(.3, 1/.3), p=1, same_on_batch=True),
)
#By setting same_on_batch=True you can apply the same transform across the batch
""",
f"""
nn.Sequential(
kornia.augmentation.RandomResizedCrop(size=({scaler}, {scaler}), scale=(3., 3.), ratio=(2., 2.), p=1.),
kornia.augmentation.RandomHorizontalFlip(p=0.7),
kornia.augmentation.RandomGrayscale(p=0.5),
)
"""
]
selected_transform = st.selectbox(
"Pick an augmentation pipeline example:", predefined_transforms
)
st.write("Transform to apply:")
readonly = False
content = st_ace(
value=selected_transform,
height=150,
language="python",
keybinding="vscode",
show_gutter=True,
show_print_margin=True,
wrap=False,
auto_update=False,
readonly=readonly,
)
if content:
transform = set_transform(content)
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)
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
cols[i].image(F.to_pil_image(x), use_column_width=True)
st.markdown(
"There are a lot more transformations available: [Documentation](https://kornia.readthedocs.io/en/latest/augmentation.module.html)"
)
st.markdown(
"Kornia can do a lot more than augmentations~ [Check it out](https://kornia.readthedocs.io/en/latest/get-started/introduction.html)"
) |