Spaces:
Sleeping
Sleeping
File size: 2,055 Bytes
248b38d ffec207 6ddafa8 3e6e09c 248b38d 961f6bd a385b9c 961f6bd a6ff839 d9e2f7d 092ff29 a6ff839 d9e2f7d 092ff29 d9e2f7d 6ddafa8 d9e2f7d a6ff839 db0d96f a385b9c db0d96f 43b044b db0d96f a385b9c 092ff29 db0d96f 7bfd7cf db0d96f a385b9c 961f6bd a6ff839 248b38d |
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 |
import streamlit as st
import torch
from transformers import AutoProcessor, UdopForConditionalGeneration
from PIL import Image
# from datasets import load_dataset
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
processor = AutoProcessor.from_pretrained("microsoft/udop-large", apply_ocr=True)
model = UdopForConditionalGeneration.from_pretrained("microsoft/udop-large")
st.title("CIC Demo (by ITT)")
st.write("Upload and Select a document (/an image) to test the model.")
# File selection
uploaded_files = st.file_uploader("Upload document(s) [/image(s)]:", type=["docx", "pdf", "pptx", "jpg", "jpeg", "png"], accept_multiple_files=True)
selected_file = st.selectbox("Select a document (/an image):", uploaded_files, format_func=lambda file: file.name if file else "None")
# Display selected file
if selected_file is not None and selected_file != "None":
file_extension = selected_file.name.split(".")[-1]
if file_extension in ["jpg", "jpeg", "png"]:
image = Image.open(selected_file).convert("RGB")
st.image(selected_file, caption="Selected Image")
else:
st.write("Selected file: ", selected_file.name)
# Model Testing
## Test button
testButton = st.button("Test Model")
## Question (/Prompt)
# question = "Question answering. How many unsafe practice of Lifting Operation?"
question = "Question answering. Is is a scene of Lifting Operation?"
## Perform Model Testing when Image is uploaded and selected as well as Test button is pressed
if testButton and selected_file != "None":
st.write("Testing the model with the selected image...")
# encoding = processor(image, question, words, boxes=boxes, return_tensors="pt")
model_encoding = processor(images=image, text=question, return_tensors="pt")
model_output = model.generate(**model_encoding)
output_text = processor.batch_decode(model_output, skip_special_tokens=True)[0]
st.write(output_text)
elif testButton and selected_file == "None":
st.write("Please upload and select a document (/an image).")
|