Spaces:
Running
on
CPU Upgrade
Running
on
CPU Upgrade
File size: 4,012 Bytes
d72ec76 |
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 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 |
import streamlit as st
import kornia
from torch import nn
import torch
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(suppress_st_warning=True)
def set_transform(content):
# st.write("set transform")
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")
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:
# 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)
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
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/introduction.html#highlighted-features)"
)
# if process:
# pass
|