feat: new flow and new Unstructured receipt parser

app.py

@@ -68,7 +68,7 @@ photo = None
 with st.sidebar:
     information = st.radio(
     "What information inside the 🧾s are you interested in extracting?",
-    ('Receipt Summary', 'Receipt Menu Details', 'Extract all'))
+    ('Receipt Summary', 'Receipt Menu Details', 'Extract all', 'Unstructured.io Parser'))
     receipt = st.selectbox('Pick one 🧾', ['1', '2', '3', '4', '5', '6'], index=1)
 
     # file upload
@@ -103,44 +103,51 @@ else:
 with col1:
     st.image(image, caption='Your target receipt')
 
-with st.spinner(f'baking the 🍩s...'):
-    if information == 'Receipt Summary':
-        processor = DonutProcessor.from_pretrained("unstructuredio/donut-base-sroie")
-        pretrained_model = VisionEncoderDecoderModel.from_pretrained("unstructuredio/donut-base-sroie")
-        task_prompt = f"<s>"
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        pretrained_model.to(device)
-    
-    elif information == 'Receipt Menu Details':
-        processor = DonutProcessor.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")
-        pretrained_model = VisionEncoderDecoderModel.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")
-        task_prompt = f"<s_cord-v2>"
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        pretrained_model.to(device)
+if st.button('Parse receipt! 🐍'):
+    with st.spinner(f'baking the 🍩s...'):
+        if information == 'Receipt Summary':
+            processor = DonutProcessor.from_pretrained("unstructuredio/donut-base-sroie")
+            pretrained_model = VisionEncoderDecoderModel.from_pretrained("unstructuredio/donut-base-sroie")
+            task_prompt = f"<s>"
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            pretrained_model.to(device)
         
-    else:
-        processor_a = DonutProcessor.from_pretrained("unstructuredio/donut-base-sroie")
-        processor_b = DonutProcessor.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")
-        pretrained_model_a = VisionEncoderDecoderModel.from_pretrained("unstructuredio/donut-base-sroie")
-        pretrained_model_b = VisionEncoderDecoderModel.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")
+        elif information == 'Receipt Menu Details':
+            processor = DonutProcessor.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")
+            pretrained_model = VisionEncoderDecoderModel.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")
+            task_prompt = f"<s_cord-v2>"
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            pretrained_model.to(device)
+            
+        elif information == 'Unstructured.io Parser':
+            processor = DonutProcessor.from_pretrained("unstructuredio/donut-base-labelstudio-A1.0")
+            pretrained_model = VisionEncoderDecoderModel.from_pretrained("unstructuredio/donut-base-labelstudio-A1.0")
+            task_prompt = f"<s>"
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            pretrained_model.to(device)
+            
+        else:  # Extract all
+            processor_a = DonutProcessor.from_pretrained("unstructuredio/donut-base-sroie")
+            processor_b = DonutProcessor.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")
+            pretrained_model_a = VisionEncoderDecoderModel.from_pretrained("unstructuredio/donut-base-sroie")
+            pretrained_model_b = VisionEncoderDecoderModel.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")
+            device = "cuda" if torch.cuda.is_available() else "cpu"
         
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-
-with col2:
-    if information == 'Extract all':
-        st.info(f'parsing 🧾 (extracting all)...')
-        pretrained_model, processor, task_prompt = pretrained_model_a, processor_a, f"<s>"
-        pretrained_model.to(device)
-        parsed_receipt_info_a, _ = run_prediction(image)
-        pretrained_model, processor, task_prompt = pretrained_model_b, processor_b, f"<s_cord-v2>"
-        pretrained_model.to(device)
-        parsed_receipt_info_b, _ = run_prediction(image)
-        st.text(f'\nReceipt Summary:')
-        st.json(parsed_receipt_info_a)
-        st.text(f'\nReceipt Menu Details:')
-        st.json(parsed_receipt_info_b)
-    else:
-        st.info(f'parsing 🧾...')
-        parsed_receipt_info, _ = run_prediction(image)
-        st.text(f'\n{information}')
-        st.json(parsed_receipt_info)
+    with col2:
+        if information == 'Extract all':
+            st.info(f'parsing 🧾 (extracting all)...')
+            pretrained_model, processor, task_prompt = pretrained_model_a, processor_a, f"<s>"
+            pretrained_model.to(device)
+            parsed_receipt_info_a, _ = run_prediction(image)
+            pretrained_model, processor, task_prompt = pretrained_model_b, processor_b, f"<s_cord-v2>"
+            pretrained_model.to(device)
+            parsed_receipt_info_b, _ = run_prediction(image)
+            st.text(f'\nReceipt Summary:')
+            st.json(parsed_receipt_info_a)
+            st.text(f'\nReceipt Menu Details:')
+            st.json(parsed_receipt_info_b)
+        else:
+            st.info(f'parsing 🧾...')
+            parsed_receipt_info, _ = run_prediction(image)
+            st.text(f'\n{information}')
+            st.json(parsed_receipt_info)