Spaces:
Sleeping
Sleeping
| import numpy as np | |
| import re | |
| import streamlit as st | |
| import torch | |
| from transformers import AutoProcessor, UdopForConditionalGeneration | |
| from PIL import Image, ImageDraw | |
| # from datasets import load_dataset | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| # UDOP uses 501 special loc ("location") tokens | |
| LAYOUT_VOCAB_SIZE = 501 | |
| def extract_coordinates(string): | |
| # Using regular expression to find all numbers in the string | |
| numbers = re.findall(r'\d+', string) | |
| # Converting the numbers to integers | |
| numbers = list(map(int, numbers)) | |
| # Ensuring there are exactly 4 numbers | |
| if len(numbers) >= 4: #if len(numbers) != 4: | |
| numbers = numbers[-4:] | |
| # Extracting coordinates | |
| x1, y1, x2, y2 = numbers | |
| else: | |
| return [] | |
| return [x1, y1, x2, y2] | |
| def unnormalize_box(box, image_width, image_height): | |
| x1 = box[0] / LAYOUT_VOCAB_SIZE * image_width | |
| y1 = box[1] / LAYOUT_VOCAB_SIZE * image_height | |
| x2 = box[2] / LAYOUT_VOCAB_SIZE * image_width | |
| y2 = box[3] / LAYOUT_VOCAB_SIZE * image_height | |
| return [x1, y1, x2, y2] | |
| processor = AutoProcessor.from_pretrained("microsoft/udop-large", apply_ocr=True) | |
| model = UdopForConditionalGeneration.from_pretrained("microsoft/udop-large") | |
| st.title("GenAI Demo (by ITT)") | |
| st.text("Upload and Select a document (/an image) to test the model.") | |
| #2 column layout | |
| col1, col2 = st.columns(2) | |
| with col1: | |
| # File selection | |
| uploaded_files = st.file_uploader("Upload document(s) [/image(s)]:", type=["docx", "pdf", "pptx", "jpg", "jpeg", "png"], accept_multiple_files=True, key="fileUpload") | |
| selected_file = st.selectbox("Select a document (/an image):", uploaded_files, format_func=lambda file: file.name if file else "None", key="fileSelect") | |
| # 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 | |
| with col2: | |
| ## Question (/Prompt) | |
| # question = "Question answering. How many unsafe practice of Lifting Operation?" | |
| default_question = "Is this a Lifting Operation scene?" | |
| task_type = st.selectbox("Question Type:", ("Classification", "Question Answering", "Layout Analysis"), index=1, key="taskSelect") | |
| question_text = st.text_area("Prompt:", placeholder=default_question, key="questionInput") | |
| if question_text is not None: | |
| question = task_type + ". " + question_text | |
| else: | |
| question = task_type + ". " + default_question | |
| ## Test button | |
| testButton = st.button("Test Model", key="testStart") | |
| ## 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) | |
| match task_type: | |
| case "Classification": | |
| output_text = processor.batch_decode(model_output, skip_special_tokens=True)[0] | |
| st.write(output_text) | |
| case "Question Answering": | |
| output_text = processor.batch_decode(model_output, skip_special_tokens=True)[0] | |
| st.write(output_text) | |
| case "Layout Analysis": | |
| output_text = processor.batch_decode(model_output, skip_special_tokens=False)[0] | |
| mean = processor.image_processor.image_mean | |
| std = processor.image_processor.image_std | |
| unnormalized_image = (model_encoding.pixel_values.squeeze().numpy() * np.array(std)[:, None, None]) + np.array(mean)[:, None, None] | |
| unnormalized_image = (unnormalized_image * 255).astype(np.uint8) | |
| unnormalized_image = np.moveaxis(unnormalized_image, 0, -1) | |
| unnormalized_image = Image.fromarray(unnormalized_image) | |
| # Get the coordinates from the output text and denormalize them | |
| coordinates = extract_coordinates(output_text) | |
| if coordinates: | |
| coordinates = unnormalize_box(coordinates, unnormalized_image.width, unnormalized_image.height) | |
| draw = ImageDraw.Draw(unnormalized_image) | |
| draw.rectangle(coordinates, outline="red") | |
| st.image(unnormalized_image, caption="Output Image") | |
| else: | |
| st.write("Cannot obtain Bounding Box coordinates: " + output_text) | |
| elif testButton and selected_file == "None": | |
| st.write("Please upload and select a document (/an image).") | |