Adding time

app.py

@@ -12,6 +12,7 @@ from docling_core.types.doc.document import DocTagsDocument
 from transformers import AutoProcessor, AutoModelForVision2Seq
 from transformers.image_utils import load_image
 from pathlib import Path
+import time
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -36,34 +37,35 @@ def get_pdf_page_count(pdf_path):
     return len(reader.pages)
 
 def get_page_image(pdf_path, page_num):
+    start = time.time()
     images = convert_from_path(pdf_path, first_page=page_num, last_page=page_num)
     page_image = images[0]
-    return page_image
+    runtime = time.time() - start
+    return page_image, f"{runtime:.2f} s"
 
 def get_docling_ocr(pdf_path, page_num):
+    start = time.time()
     result = converter.convert(pdf_path, page_range=(page_num, page_num))
     markdown_text_docling = result.document.export_to_markdown()
-    return markdown_text_docling
+    runtime = time.time() - start
+    return markdown_text_docling, f"{runtime:.2f} s"
 
 def get_paddle_ocr(pdf_path, page_num):
-    page_image = get_page_image(pdf_path, page_num)
-
+    start = time.time()
+    page_image = get_page_image(pdf_path, page_num)[0]
     output = pipeline.predict(input=np.array(page_image))
-
     markdown_list = []
-
     for res in output:
         md_info = res.markdown
         markdown_list.append(md_info)
-
     markdown_text_paddleOCR = pipeline.concatenate_markdown_pages(markdown_list)
-    return markdown_text_paddleOCR
+    runtime = time.time() - start
+    return markdown_text_paddleOCR, f"{runtime:.2f} s"
 
 def get_smoldocling_ocr(pdf_path, page_num):
-    page_image = get_page_image(pdf_path, page_num)
+    start = time.time()
+    page_image = get_page_image(pdf_path, page_num)[0]
     image = load_image(page_image)
-
-    # Create input messages
     messages = [
         {
             "role": "user",
@@ -73,12 +75,9 @@ def get_smoldocling_ocr(pdf_path, page_num):
             ]
         },
     ]
-
-    # Prepare inputs
     prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
     inputs = processor(text=prompt, images=[image], return_tensors="pt")
     inputs = inputs.to(DEVICE)
-
     generated_ids = model.generate(**inputs, max_new_tokens=8192)
     prompt_length = inputs.input_ids.shape[1]
     trimmed_generated_ids = generated_ids[:, prompt_length:]
@@ -86,13 +85,11 @@ def get_smoldocling_ocr(pdf_path, page_num):
         trimmed_generated_ids,
         skip_special_tokens=False,
     )[0].lstrip()
-
-    # Populate document
     doctags_doc = DocTagsDocument.from_doctags_and_image_pairs([doctags], [image])
     doc = DoclingDocument.load_from_doctags(doctags_doc, document_name="Document")
-
     markdown_text_smoldocling = doc.export_to_markdown()
-    return markdown_text_smoldocling
+    runtime = time.time() - start
+    return markdown_text_smoldocling, f"{runtime:.2f} s"
 
 title = "OCR Arena"
 description = "A simple Gradio interface to extract text from PDFs and compare OCR models"
@@ -117,16 +114,20 @@ with gr.Blocks(theme=gr.themes.Glass()) as demo:
                         clear_btn = gr.ClearButton(components=[pdf, page_num])
                         submit_btn = gr.Button("Submit", variant='primary')
 
-                    submit_btn.click(get_page_image, inputs=[pdf, page_num], outputs=original).then(    
-                        get_docling_ocr, inputs=[pdf, page_num], outputs=docling_ocr_out).then(
-                        get_paddle_ocr, inputs=[pdf, page_num], outputs=paddle_ocr_out).then(
-                        get_smoldocling_ocr, inputs=[pdf, page_num], outputs=smoldocling_ocr_out)
+                    submit_btn.click(get_page_image, inputs=[pdf, page_num], outputs=[original, original_runtime]).then(    
+                        get_docling_ocr, inputs=[pdf, page_num], outputs=[docling_ocr_out, docling_ocr_runtime]).then(
+                        get_paddle_ocr, inputs=[pdf, page_num], outputs=[paddle_ocr_out, paddle_ocr_runtime]).then(
+                        get_smoldocling_ocr, inputs=[pdf, page_num], outputs=[smoldocling_ocr_out, smoldocling_ocr_runtime])
 
         with gr.Column():
             original = gr.Image(width=640, height=640, label="Original Page", interactive=False)
+            original_runtime = gr.Textbox(label="Image Extraction Time", type="text", interactive=False)
             docling_ocr_out = gr.Textbox(label="Docling OCR Output", type="text", interactive=False)
+            docling_ocr_runtime = gr.Textbox(label="Docling OCR Time", type="text", interactive=False)
             paddle_ocr_out = gr.Textbox(label="Paddle OCR Output", type="text", interactive=False)
+            paddle_ocr_runtime = gr.Textbox(label="Paddle OCR Time", type="text", interactive=False)
             smoldocling_ocr_out = gr.Textbox(label="SmolDocling OCR Output", type="text", interactive=False)
+            smoldocling_ocr_runtime = gr.Textbox(label="SmolDocling OCR Time", type="text", interactive=False)
 
     examples_obj = gr.Examples(examples=examples, inputs=[pdf])