Spaces:
Runtime error
Runtime error
import streamlit as st | |
import io | |
import sys | |
import time | |
import json | |
sys.path.append("./virtex/") | |
def gen_show_caption(sub_prompt=None, cap_prompt = ""): | |
with st.spinner("Generating Caption"): | |
if sub_prompt is None and cap_prompt is not "": | |
st.write("Without a specified subreddit, caption prompts will skip subreddit prediction") | |
subreddit, caption = virtexModel.predict(image_dict, sub_prompt=sub_prompt, prompt = cap_prompt) | |
st.header("Predicted Caption:\n\n") | |
st.subheader(f"r/{subreddit}:\t{caption}\n") | |
st.title("Image Captioning Demo from Redcaps") | |
st.sidebar.markdown( | |
""" | |
Image Captioning Model from VirTex trained on Redcaps | |
""" | |
) | |
with st.spinner("Loading Model"): | |
from model import * | |
sample_images = get_samples() | |
virtexModel = VirTexModel() | |
imageLoader = ImageLoader() | |
valid_subs = json.load(open(VALID_SUBREDDITS_PATH)) | |
valid_subs.insert(0, None) | |
random_image = get_rand_img(sample_images) | |
rand_idx = 0 | |
st.sidebar.title("Select a sample image") | |
if st.sidebar.button("Random Sample Image"): | |
rand_idx, random_image = get_rand_img(sample_images) | |
sample_image = None | |
sample_image = st.sidebar.selectbox( | |
"", | |
sample_images, | |
index=rand_idx | |
) | |
st.sidebar.title("Select a Subreddit") | |
sub = st.sidebar.selectbox( | |
"Select None for a Predicted Subreddit", | |
valid_subs | |
) | |
st.sidebar.title("Write a Custom Prompt") | |
cap_prompt = st.sidebar.text_input( | |
"Leave this blank for an unbiased caption", | |
value="" | |
) | |
uploaded_image = None | |
with st.sidebar.form("file-uploader-form", clear_on_submit=True): | |
uploaded_file = st.file_uploader("Choose a file") | |
submitted = st.form_submit_button("Submit") | |
if uploaded_file is not None and submitted: | |
uploaded_image = Image.open(io.BytesIO(uploaded_file.getvalue())) | |
if uploaded_image is None and submitted: | |
st.write("Please select a file to upload") | |
else: | |
image_file = sample_image if sample_image is not None else random_image | |
image = uploaded_image if uploaded_image is not None else Image.open(image_file) | |
image_dict = imageLoader.transform(image) | |
# image = imageLoader.to_image(image_dict["image"].squeeze(0)) | |
show = st.image(image) | |
show.image(image, "Your Image") | |
gen_show_caption(sub, cap_prompt) | |
image.close() | |
# from model import * | |
# download_files() | |
# sample_images = get_samples() | |
# v, il = VirTexModel(), ImageLoader() | |
# for s in sample_images: | |
# subreddit, caption = v.predict(il.load(s)) | |
# print("=====================") | |
# print(subreddit) | |
# print(caption) | |