Update app.py

app.py

@@ -3,11 +3,15 @@ import torch
 from transformers import AutoProcessor, UdopForConditionalGeneration
 # from datasets import load_dataset
 
-# processor = AutoProcessor.from_pretrained("microsoft/udop-large", apply_ocr=True)
-# model = UdopForConditionalGeneration.from_pretrained("microsoft/udop-large")
+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")
+
+question = "Question answering. How many unsafe practice of Lifting Operation"
 
 st.title("CIC Demo (by ITT)")
-st.write("Select or upload a document (/an image) to test the model.")
+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)
@@ -25,10 +29,10 @@ if selected_file is not None and selected_file != "None":
 testButton = st.button("Test Model")
 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")
+    predicted_ids = model.generate(**encoding)
+    print(processor.batch_decode(predicted_ids, skip_special_tokens=True)[0])    
 elif testButton and selected_file == "None":
-    st.write("Please upload and select an image.")
+    st.write("Please upload and select a document (/an image).")
 
-# encoding = processor(image, question, words, boxes=boxes, return_tensors="pt")
 
-# predicted_ids = model.generate(**encoding)
-# print(processor.batch_decode(predicted_ids, skip_special_tokens=True)[0])