File size: 1,678 Bytes
16906c1 8d21e52 16906c1 f0fcc84 16906c1 |
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 |
import streamlit as st
from streamlit_drawable_canvas import st_canvas
import os
import utils
from PIL import Image
st.set_page_config("VAE MNIST Pytorch Lightning")
st.title("VAE Playground")
# title_img = Image.open("images/title_img.jpg")
# st.image(title_img)
st.markdown(
"This is a simple streamlit app to showcase how a simple VAEs."
)
def load_model_files():
files = os.listdir("./models/")
# Docker creates some whiteout files which mig
files = [i for i in files if ".ckpt" in i]
clean_names = [utils.parse_model_file_name(name) for name in files]
return {k: v for k, v in zip(clean_names, files)}
file_name_map = load_model_files()
files = list(file_name_map.keys())
st.header("🖼️ Image Reconstruction", "recon")
with st.form("reconstruction"):
model_name = st.selectbox("Choose Model:", files,
key="recon_model_select") or "conv_vae"
recon_model_name = file_name_map[model_name]
recon_canvas = st_canvas(
# Fixed fill color with some opacity
fill_color="rgba(255, 165, 0, 0.3)",
stroke_width=8,
stroke_color="#FFFFFF",
background_color="#000000",
update_streamlit=True,
height=150,
width=150,
drawing_mode="freedraw",
key="recon_canvas",
)
submit = st.form_submit_button("Perform Reconstruction")
if submit:
recon_model = utils.load_model(recon_model_name)
inp_tens = utils.canvas_to_tensor(recon_canvas)
_, _, out = recon_model(inp_tens)
out = (out+1)/2
out_img = utils.resize_img(utils.tensor_to_img(out), 150, 150)
if submit:
st.image(out_img)
|