Complete digitizing

app/tabs/submit.py

@@ -1,17 +1,22 @@
+import json
 import logging
 import os
+from pathlib import Path
 import time
 
-import numpy as np
-import pandas as pd
-from pathlib import Path
+from PIL import Image
 from dawsonia import io
 from dawsonia import digitize
 from dawsonia.ml import ml
-import pooch
 import gradio as gr
-import yaml
 from gradio_modal import Modal
+import numpy as np
+from numpy.typing import NDArray
+import pandas as pd
+import pooch
+import yaml
+
+from .visualizer import Page, TableCell
 
 logger = logging.getLogger(__name__)
 
@@ -39,10 +44,11 @@ if os.environ.get("GRADIO_CACHE_DIR", GRADIO_CACHE) != GRADIO_CACHE:
 
 
 def run_dawsonia(
-    table_fmt_config_override, batch_image_gallery, book, progress=gr.Progress()
+    table_fmt_config_override, first_page, last_page, book, progress=gr.Progress()
 ):
-    if None in (batch_image_gallery, book) or len(batch_image_gallery) == 0:
+    if book is None:
         raise ValueError("You need to select / upload the pages to digitize")
+
     progress(0, desc="Dawsonia: starting")
 
     model_path = Path("data/models/dawsonia/2024-07-02")
@@ -51,7 +57,7 @@ def run_dawsonia(
     print("Dawsonia: digitizing", book)
     table_fmt = book.table_format
 
-    output_path_book = output_path / book.station_name / book._name
+    output_path_book = output_path / book.station_name
     output_path_book.mkdir(exist_ok=True, parents=True)
     (output_path_book / "probablities").mkdir(exist_ok=True)
 
@@ -63,30 +69,63 @@ def run_dawsonia(
         for table_idx in table_fmt.preproc.idx_tables_size_verify
     ]
 
-    for page_number in range(len(batch_image_gallery)):
+    collection = []
+
+    for page_number in range(first_page, last_page):
         output_path_page = output_path_book / str(page_number)
-        results = [
-            digitize.digitize_page_and_write_output(
-                book,
-                init_data,
-                page_number=page_number + 3,
-                date_str="2022-02-02",
-                model_path=model_path,
-                model_predict=ml.model_predict,
-                prob_thresh=0.5,
-                output_path_page=output_path_page,
-                output_text_fmt=True,
-                debug=True,
-            )
-        ]
+        gr.Info(f"Digitizing {page_number = }")
+
+        *_, stats = digitize.digitize_page_and_write_output(
+            book,
+            init_data,
+            page_number=page_number,
+            date_str=f"0000-page-{page_number}",
+            model_path=model_path,
+            model_predict=ml.model_predict,
+            prob_thresh=0.5,
+            output_path_page=output_path_page,
+            output_text_fmt=False,
+            debug=False,
+        )
+        progress_value = (page_number - first_page) / max(1, last_page - first_page)
+        progress(progress_value, desc=f"Dawsonia: {stats!s:.50}")
+
+        collection.append(read_page(stats, output_path_book, str(page_number)))
 
-    collection = []
-    time.sleep(1)
     gr.Info("Pages were succesfully digitized ✨")
 
     yield collection, gr.skip()
 
 
+def read_page(stats: digitize.Statistics, output_path_book: Path, prefix: str):
+    if stats.tables_detected > 0:
+        values_df = pd.read_parquet((output_path_book / prefix).with_suffix(".parquet"))
+        table_meta = json.loads(
+            (output_path_book / "table_meta" / prefix).with_suffix(".json").read_text()
+        )
+        with Image.open(
+            (output_path_book / "pages" / prefix).with_suffix(".webp")
+        ) as im:
+            width = im.width
+            height = im.height
+
+        values_array = values_df.values.flatten()
+        bbox_array = np.array(table_meta["table_positions"]).reshape(
+            values_array.size, 4
+        )
+        cells = [
+            make_cell(value, bbox) for value, bbox in zip(values_array, bbox_array)
+        ]
+        return Page(width, height, cells)
+
+
+def make_cell(value: str, bbox: NDArray[np.int64]):
+    x, y, w, h = bbox
+    xmax, ymax = x+w, y+h
+    polygon = (x,y), (xmax, y), (xmax, ymax), (x, ymax), (x,y) 
+    return TableCell(polygon, text_x=x, text_y=y, text=value)
+
+
 def all_example_images() -> list[str]:
     """
     Get paths to all example images.
@@ -106,6 +145,9 @@ def get_selected_example_image(
     # for name, details in PIPELINES.items():
     name, _ext = event.value["image"]["orig_name"].split(".")
 
+    if (last_page - first_page) > MAX_IMAGES:
+        raise ValueError(f"Maximum images you can digitize is set to: {MAX_IMAGES}")
+
     if name in PIPELINES:
         book_path = pooch.retrieve(**PIPELINES[name], path=DATA_CACHE)
         first, last, book = io.read_book(book_path)
@@ -114,55 +156,6 @@ def get_selected_example_image(
         return [book.read_image(pg) for pg in range(first_page, last_page)], book
 
 
-table_fmt_config_override_placeholder = (
-    """\
-[default]
-version = 0
-
-# Default values, but wrote explicitly here. See PreprocConfig class
-[default.preproc]
-table_modif = true
-corr_rotate = true
-row_idx_unit = "HOURS"
-idx_tables_size_verify = [0, 1]
-
-[version.0]
-columns = [
-  [
-    "term_på_baro",
-    "barom",
-    "torra_term",
-    "våta_term",
-    "moln_slag_lägre",
-    "moln_mängd_lägre",
-    "moln_slag_medel",
-    "moln_slag_högre"
-  ],
-  [
-    "moln_het_sol_dimma_nederbörd_total",
-    "vind_riktning",
-    "vind_beaufort",
-    "vind_m_sek",
-    "sikt",
-    "sjögang",
-    "maximi_term",
-    "minimi_term",
-    "nederbörd_mängd",
-    "nederbörd_slag"
-  ]
-]
-name_idx = "tid"
-rows = [2, 8, 14, 19, 21]
-tables = [
-  [5, 8],
-  [5, 10],
-  [3, 1],
-  [4, 2],
-  [4, 5]
-]
-""",
-)
-
 with gr.Blocks() as submit:
     gr.Markdown("# Upload")
     gr.Markdown(
@@ -180,7 +173,7 @@ with gr.Blocks() as submit:
                     label="Image to digitize",
                     interactive=True,
                     object_fit="scale-down",
-                    scale=10,
+                    scale=1,
                 )
 
             with gr.Column(scale=2):
@@ -225,6 +218,7 @@ with gr.Blocks() as submit:
         get_selected_example_image,
         (first_page, last_page),
         (batch_image_gallery, batch_book_state),
+        trigger_mode="always_last",
     )
 
     @batch_image_gallery.upload(
@@ -239,7 +233,7 @@ with gr.Blocks() as submit:
 
     run_button.click(
         fn=run_dawsonia,
-        inputs=[table_fmt_config_override, batch_image_gallery, batch_book_state],
-        outputs=[collection_submit_state, batch_image_gallery],
+        inputs=(table_fmt_config_override, first_page, last_page, batch_book_state),
+        outputs=(collection_submit_state, batch_image_gallery),
     )
     edit_table_fmt_button.click(lambda: Modal(visible=True), None, edit_table_fmt_modal)