progress thingy

app.py

@@ -4,15 +4,17 @@ import os
 from pathlib import Path
 from queue import SimpleQueue
 from threading import Thread
+from time import sleep
 from typing import Any
 
 import gradio as gr  # type: ignore
 import rerun as rr
+import rerun.blueprint as rrb
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from gradio_rerun import Rerun  # type: ignore
 
-from ocr import detect_and_log_layouts
+from ocr import detect_and_log_layouts, PAGE_LIMIT
 
 CUSTOM_PATH = "/"
 
@@ -27,18 +29,42 @@ app.add_middleware(
     allow_origins=origins,
 )
 
-def file_ocr(log_queue: SimpleQueue[Any], file_path: str):
-    detect_and_log_layouts(log_queue, file_path)
+
+def progress_log(log_queue: SimpleQueue[Any], done: SimpleQueue[Any]):
+    dots = 0
+    while True:
+        if not done.empty():
+            break
+        sleep(0.7)
+        log_queue.put([
+            "log",
+            "progress",
+            [rr.TextDocument(f"working{'.'*(dots+1)}")]
+        ])
+        dots = (dots + 1) % 5
+
+
+def file_ocr(log_queue: SimpleQueue[Any], file_path: str, start_page: int, end_page: int):
+    detect_and_log_layouts(log_queue, file_path, start_page, end_page)
     log_queue.put("done")
 
+
 @rr.thread_local_stream("PaddleOCR")
-def log_to_rr(file_path: Path):
+def log_to_rr(file_path: Path, start_page: int = 1, end_page: int = -1):
     stream = rr.binary_stream()
 
     log_queue: SimpleQueue[Any] = SimpleQueue()
-    handle = Thread(target=file_ocr, args=[log_queue, str(file_path)])
+    done: SimpleQueue[Any] = SimpleQueue()
+    Thread(target=progress_log, args=[log_queue, done]).start()
+    handle = Thread(target=file_ocr, args=[log_queue, str(file_path), start_page, end_page])
     handle.start()
 
+    rr.send_blueprint(rrb.Blueprint(
+        rrb.TextDocumentView(contents=["progress/**"]),
+        collapse_panels=True,
+    ))
+    yield stream.read()
+
     while True:
         msg = log_queue.get()
         if msg == "done":
@@ -53,13 +79,13 @@ def log_to_rr(file_path: Path):
             entity_path = msg[1]
             args = msg[2]
             kwargs = msg[3] if len(msg) >= 4 else {}
-            # print(entity_path)
-            # print(args)
-            # print(kwargs)
             rr.log(entity_path, *args, **kwargs)
 
         yield stream.read()
 
+    rr.log("progress",rr.TextDocument("Done!"))
+    yield stream.read()
+    done.put(())
     handle.join()
     print("done")
 
@@ -73,21 +99,48 @@ with gr.Blocks() as demo:
     gr.Markdown(DESCRIPTION)
     with gr.Row():
         with gr.Column(scale=1):
-            with gr.Row():
-                #input_image = gr.Image(label="Input Image", image_mode="RGBA", sources="upload", type="filepath")
-                input_file = gr.File(label="Input file (image/pdf)")
-            with gr.Row():
-                button = gr.Button()
-            with gr.Row():
-                gr.Examples(
-                    examples=[os.path.join("examples", img_name) for img_name in sorted(os.listdir("examples"))],
-                    inputs=[input_file],
-                    label="Examples",
-                    cache_examples=False,
-                    examples_per_page=12,
-                )
+            with gr.Tab(label="Upload Image"):
+                with gr.Row():
+                    input_image_file = gr.Image(label="Input Image", image_mode="RGBA", sources="upload", type="filepath")
+                    # input_image_file = gr.Image(label="Input image")
+                with gr.Row():
+                    image_button = gr.Button()
+                with gr.Row():
+                    gr.Examples(
+                        examples=[
+                            os.path.join("image_examples", img_name)
+                            for img_name in sorted(os.listdir("image_examples"))
+                        ],
+                        inputs=[input_image_file],
+                        label="Examples",
+                        cache_examples=False,
+                        examples_per_page=12,
+                    )
+            with gr.Tab(label="Upload pdf"):
+                with gr.Row():
+                    input_pdf_file = gr.File(label="Input pdf")
+                gr.Markdown(f"Max {PAGE_LIMIT} pages, -1 on end page means max number of pages")
+                with gr.Row():
+                    start_page_number = gr.Number(1, label="Start page", minimum=1)
+                with gr.Row():
+                    end_page_number = gr.Number(-1, label="End page")
+                with gr.Row():
+                    pdf_button = gr.Button()
+                with gr.Row():
+                    gr.Examples(
+                        examples=[
+                            os.path.join("pdf_examples", img_name) for img_name in sorted(os.listdir("pdf_examples"))
+                        ],
+                        inputs=[input_pdf_file],
+                        label="Examples",
+                        cache_examples=False,
+                        examples_per_page=12,
+                    )
         with gr.Column(scale=4):
             viewer = Rerun(streaming=True, height=900)
-    button.click(log_to_rr, inputs=[input_file], outputs=[viewer])
+
+        image_button.click(log_to_rr, inputs=[input_image_file], outputs=[viewer])
+        pdf_button.click(log_to_rr, inputs=[input_pdf_file, start_page_number, end_page_number], outputs=[viewer])
+
 
 app = gr.mount_gradio_app(app, demo, path=CUSTOM_PATH)

ocr.py

@@ -25,6 +25,8 @@ DATASET_DIR: Final = EXAMPLE_DIR / "dataset"
 
 SAMPLE_IMAGE_URLs = ["https://storage.googleapis.com/rerun-example-datasets/ocr/paper.png"]
 
+PAGE_LIMIT = 10
+
 LayoutStructure: TypeAlias = tuple[
     list[str], list[str], list[rrb.Spatial2DView], list[rrb.Spatial2DView], list[rrb.Spatial2DView]
 ]
@@ -352,7 +354,11 @@ def generate_blueprint(
                         contents=[f"{page_path}/Image/**"] + detections_paths,
                     ),
                     rrb.Spatial2DView(name="Detections", contents=[f"{page_path}/Image/**"]),
-                    rrb.TextDocumentView(name="Recovery", contents=f"{page_path}/Recovery"),
+                    rrb.Vertical(
+                        rrb.TextDocumentView(name="Progress", contents=["progress/**"]),
+                        rrb.TextDocumentView(name="Recovery", contents=f"{page_path}/Recovery"),
+                        row_shares=[1, 4]
+                    )
                 ),
                 rrb.Horizontal(*section_tabs),
                 name=page_path,
@@ -366,11 +372,22 @@ def generate_blueprint(
     )
 
 
-def detect_and_log_layouts(log_queue: SimpleQueue[Any], file_path: str) -> None:
+def detect_and_log_layouts(log_queue: SimpleQueue[Any], file_path: str, start_page: int = 1, end_page: int | None = -1) -> None:
+    if end_page == -1:
+        end_page = start_page + PAGE_LIMIT
+    if end_page < start_page:
+        end_page = start_page
+
     images: list[npt.NDArray[np.uint8]] = []
     if file_path.endswith(".pdf"):
         # convert pdf to images
-        images.extend(np.array(img, dtype=np.uint8) for img in pdf2image.convert_from_path(file_path))
+        images.extend(np.array(img, dtype=np.uint8) for img in pdf2image.convert_from_path(file_path, first_page=start_page, last_page=end_page))
+        if len(images) > PAGE_LIMIT:
+            log_queue.put([
+                "log",
+                "error",
+                [rr.TextLog(f"Too many pages requsted: {len(images)} requested but the limit is {PAGE_LIMIT}")],
+            ])
     else:
         # read image
         img = cv2.imread(file_path)
@@ -379,7 +396,7 @@ def detect_and_log_layouts(log_queue: SimpleQueue[Any], file_path: str) -> None:
 
     # Extracte the layout from each image
     layouts: list[Layout] = []
-    page_paths = [f"page_{i + 1}" for i in range(len(images))]
+    page_paths = [f"page_{i + start_page}" for i in range(len(images))]
     processed_layouts: list[LayoutStructure] = []
     for i, (image, page_path) in enumerate(zip(images, page_paths)):
         layouts.append(detect_and_log_layout(log_queue, image, page_path))