Ashwin V. Mohanan commited on
Commit
5015c53
·
1 Parent(s): ed479a9

Variable probability threshold

Browse files
Files changed (2) hide show
  1. app/tabs/submit.py +94 -31
  2. app/tabs/visualizer.py +5 -4
app/tabs/submit.py CHANGED
@@ -9,6 +9,7 @@ from PIL import Image
9
  from dawsonia import io
10
  from dawsonia import digitize
11
  from dawsonia.ml import ml
 
12
  import gradio as gr
13
  from gradio_modal import Modal
14
  import numpy as np
@@ -40,16 +41,22 @@ PIPELINES: dict[str, dict[str, str]] = {
40
  "härnösand": dict(
41
  url="https://git.smhi.se/ai-for-obs/data/-/raw/688c04f13e8e946962792fe4b4e0ded98800b154/raw_zarr/H%C3%84RN%C3%96SAND/DAGBOK_H%C3%A4rn%C3%B6sand_Station_1934.zarr.zip",
42
  known_hash="sha256:a58fdb6521214d0bd569c9325ce78d696738de28ce6ec869cde0d46616b697f2",
43
- )
44
  }
45
 
46
 
47
  def run_dawsonia(
48
- table_fmt_config_override, first_page, last_page, book, gallery, progress=gr.Progress()
 
 
 
 
 
 
49
  ):
50
  if book is None:
51
  raise ValueError("You need to select / upload the pages to digitize")
52
-
53
  progress(0, desc="Dawsonia: starting")
54
 
55
  model_path = Path("data/models/dawsonia/2024-07-02")
@@ -79,7 +86,11 @@ def run_dawsonia(
79
  output_path_page = output_path_book / str(page_number)
80
  gr.Info(f"Digitizing {page_number = }")
81
 
82
- if not (output_path_book / str(page_number)).with_suffix(".parquet").exists():
 
 
 
 
83
  digitize.digitize_page_and_write_output(
84
  book,
85
  init_data,
@@ -87,48 +98,66 @@ def run_dawsonia(
87
  date_str=f"0000-page-{page_number}",
88
  model_path=model_path,
89
  model_predict=ml.model_predict,
90
- prob_thresh=0.5,
91
  output_path_page=output_path_page,
92
  output_text_fmt=False,
93
  debug=False,
94
  )
95
  progress_value = (page_number - first_page) / max(1, last_page - first_page)
96
 
97
- if results := read_page(output_path_book, str(page_number), progress, progress_value): # , im_from_gallery[0])
 
 
 
 
 
 
98
  page, im = results
99
  collection.append(page)
100
  images.append(im)
101
  else:
102
  gr.Info(f"No tables detected in {page_number = }")
103
 
104
-
105
  gr.Info("Pages were succesfully digitized ✨")
106
 
107
  # yield collection, images
108
  yield collection, gr.skip()
109
 
110
 
111
- def read_page(output_path_book: Path, prefix: str, progress, progress_value, im_path_from_gallery: str = ""):
112
- stats = digitize.Statistics.from_json((output_path_book / "statistics" / prefix).with_suffix(".json"))
 
 
 
 
 
 
 
 
 
113
  print(stats)
114
  progress(progress_value, desc=f"Dawsonia: {stats!s:.50}")
115
  if stats.tables_detected > 0:
116
  values_df = pd.read_parquet((output_path_book / prefix).with_suffix(".parquet"))
 
 
 
117
  table_meta = json.loads(
118
  (output_path_book / "table_meta" / prefix).with_suffix(".json").read_text()
119
  )
120
  with Image.open(
121
- image_path:=(output_path_book / "pages" / prefix).with_suffix(".webp")
122
  ) as im:
123
  width = im.width
124
  height = im.height
125
 
126
  values_array = values_df.values.flatten()
127
- bbox_array = np.hstack(table_meta["table_positions"]).reshape(
128
- -1, 4
129
- )
130
  cells = [
131
- make_cell(value, bbox) for value, bbox in zip(values_array, bbox_array)
 
 
132
  ]
133
 
134
  return Page(width, height, cells, im_path_from_gallery or str(image_path)), im
@@ -136,10 +165,10 @@ def read_page(output_path_book: Path, prefix: str, progress, progress_value, im_
136
 
137
  def make_cell(value: str, bbox: NDArray[np.int64]):
138
  y, x, h, w = bbox
139
- xmin, ymin = x-w//2, y-h//2
140
- xmax, ymax = x+w//2, y+h//2
141
- polygon = (xmin,ymin), (xmax, ymin), (xmax, ymax), (xmin, ymax), (xmin,ymin)
142
- return TableCell(polygon, text_x=x, text_y=y, text=value)
143
 
144
 
145
  def all_example_images() -> list[str]:
@@ -171,7 +200,12 @@ def get_selected_example_image(
171
  first, last, book = io.read_book(book_path)
172
  book._name = name
173
  book.size_cell = [1.0, 1.0, 1.0, 1.0]
174
- return [book.read_image(pg) for pg in range(first_page, last_page)], book, book_path, station_tf.read_text()
 
 
 
 
 
175
 
176
 
177
  def overwrite_table_format_file(book: io.Book, book_path, table_fmt: str):
@@ -184,9 +218,8 @@ def overwrite_table_format_file(book: io.Book, book_path, table_fmt: str):
184
 
185
 
186
  with gr.Blocks() as submit:
187
- gr.Markdown("# Upload")
188
  gr.Markdown(
189
- "Select or upload the image you want to transcribe. You can upload up to five images at a time."
190
  )
191
 
192
  batch_book_state = gr.State()
@@ -197,11 +230,11 @@ with gr.Blocks() as submit:
197
  with gr.Row(equal_height=True):
198
  with gr.Column(scale=5):
199
  batch_image_gallery = gr.Gallery(
200
- file_types=["image"],
201
- label="Image to digitize",
202
  interactive=True,
203
  object_fit="scale-down",
204
- scale=1,
205
  )
206
 
207
  with gr.Column(scale=2):
@@ -215,6 +248,7 @@ with gr.Blocks() as submit:
215
  object_fit="scale-down",
216
  min_width=250,
217
  )
 
218
 
219
  with Modal(visible=False) as edit_table_fmt_modal:
220
  with gr.Column():
@@ -223,7 +257,9 @@ with gr.Blocks() as submit:
223
  "Write a custom table format, overriding the default one. "
224
  "Click on the **Save** button when you are done."
225
  )
226
- save_tf_button = gr.Button("Save", variant="primary", scale=0, min_width=200)
 
 
227
  gr.HTML(
228
  (
229
  "<a href='https://dawsonia.readthedocs.io/en/latest/user_guide/misc.html#table-formats' target='_blank'>"
@@ -235,6 +271,15 @@ with gr.Blocks() as submit:
235
  )
236
  table_fmt_config_override = gr.Code("", language="python")
237
 
 
 
 
 
 
 
 
 
 
238
  with gr.Row():
239
  run_button = gr.Button("Digitize", variant="primary", scale=0, min_width=200)
240
  edit_table_fmt_button = gr.Button(
@@ -246,7 +291,12 @@ with gr.Blocks() as submit:
246
  examples.select(
247
  get_selected_example_image,
248
  (first_page, last_page),
249
- (batch_image_gallery, batch_book_state, batch_book_path_state, table_fmt_config_override),
 
 
 
 
 
250
  trigger_mode="always_last",
251
  )
252
 
@@ -255,21 +305,34 @@ with gr.Blocks() as submit:
255
  outputs=[batch_image_gallery],
256
  )
257
  def validate_images(images):
 
258
  if len(images) > MAX_IMAGES:
259
  gr.Warning(f"Maximum images you can upload is set to: {MAX_IMAGES}")
260
  return gr.update(value=None)
261
-
262
- gr.Warning("Digitizing uploaded images is not implemented yet! Work in progress!")
 
 
263
  raise NotImplementedError("WIP")
264
  return images
265
-
266
  run_button.click(
267
  fn=run_dawsonia,
268
- inputs=(table_fmt_config_override, first_page, last_page, batch_book_state, batch_image_gallery),
 
 
 
 
 
 
 
269
  outputs=(collection_submit_state, batch_image_gallery),
270
  )
271
 
272
  ## Table formats modal dialog box
273
  edit_table_fmt_button.click(lambda: Modal(visible=True), None, edit_table_fmt_modal)
274
  save_tf_button.click(
275
- overwrite_table_format_file, (batch_book_state, batch_book_path_state, table_fmt_config_override), (batch_book_state,))
 
 
 
 
9
  from dawsonia import io
10
  from dawsonia import digitize
11
  from dawsonia.ml import ml
12
+ from dawsonia.typing import Probability
13
  import gradio as gr
14
  from gradio_modal import Modal
15
  import numpy as np
 
41
  "härnösand": dict(
42
  url="https://git.smhi.se/ai-for-obs/data/-/raw/688c04f13e8e946962792fe4b4e0ded98800b154/raw_zarr/H%C3%84RN%C3%96SAND/DAGBOK_H%C3%A4rn%C3%B6sand_Station_1934.zarr.zip",
43
  known_hash="sha256:a58fdb6521214d0bd569c9325ce78d696738de28ce6ec869cde0d46616b697f2",
44
+ ),
45
  }
46
 
47
 
48
  def run_dawsonia(
49
+ table_fmt_config_override,
50
+ first_page,
51
+ last_page,
52
+ prob_thresh,
53
+ book,
54
+ gallery,
55
+ progress=gr.Progress(),
56
  ):
57
  if book is None:
58
  raise ValueError("You need to select / upload the pages to digitize")
59
+
60
  progress(0, desc="Dawsonia: starting")
61
 
62
  model_path = Path("data/models/dawsonia/2024-07-02")
 
86
  output_path_page = output_path_book / str(page_number)
87
  gr.Info(f"Digitizing {page_number = }")
88
 
89
+ if (
90
+ not (output_path_book / str(page_number))
91
+ .with_suffix(".parquet")
92
+ .exists()
93
+ ):
94
  digitize.digitize_page_and_write_output(
95
  book,
96
  init_data,
 
98
  date_str=f"0000-page-{page_number}",
99
  model_path=model_path,
100
  model_predict=ml.model_predict,
101
+ prob_thresh=prob_thresh,
102
  output_path_page=output_path_page,
103
  output_text_fmt=False,
104
  debug=False,
105
  )
106
  progress_value = (page_number - first_page) / max(1, last_page - first_page)
107
 
108
+ if results := read_page(
109
+ output_path_book,
110
+ str(page_number),
111
+ prob_thresh,
112
+ progress,
113
+ progress_value,
114
+ ): # , im_from_gallery[0])
115
  page, im = results
116
  collection.append(page)
117
  images.append(im)
118
  else:
119
  gr.Info(f"No tables detected in {page_number = }")
120
 
 
121
  gr.Info("Pages were succesfully digitized ✨")
122
 
123
  # yield collection, images
124
  yield collection, gr.skip()
125
 
126
 
127
+ def read_page(
128
+ output_path_book: Path,
129
+ prefix: str,
130
+ prob_thresh: float,
131
+ progress,
132
+ progress_value,
133
+ im_path_from_gallery: str = "",
134
+ ):
135
+ stats = digitize.Statistics.from_json(
136
+ (output_path_book / "statistics" / prefix).with_suffix(".json")
137
+ )
138
  print(stats)
139
  progress(progress_value, desc=f"Dawsonia: {stats!s:.50}")
140
  if stats.tables_detected > 0:
141
  values_df = pd.read_parquet((output_path_book / prefix).with_suffix(".parquet"))
142
+ prob_df = pd.read_parquet(
143
+ (output_path_book / "probablities" / prefix).with_suffix(".parquet")
144
+ )
145
  table_meta = json.loads(
146
  (output_path_book / "table_meta" / prefix).with_suffix(".json").read_text()
147
  )
148
  with Image.open(
149
+ image_path := (output_path_book / "pages" / prefix).with_suffix(".webp")
150
  ) as im:
151
  width = im.width
152
  height = im.height
153
 
154
  values_array = values_df.values.flatten()
155
+ prob_array = prob_df.values.flatten()
156
+ bbox_array = np.hstack(table_meta["table_positions"]).reshape(-1, 4)
 
157
  cells = [
158
+ make_cell(value, bbox)
159
+ for value, prob, bbox in zip(values_array, prob_array, bbox_array)
160
+ if prob > prob_thresh
161
  ]
162
 
163
  return Page(width, height, cells, im_path_from_gallery or str(image_path)), im
 
165
 
166
  def make_cell(value: str, bbox: NDArray[np.int64]):
167
  y, x, h, w = bbox
168
+ xmin, ymin = x - w // 2, y - h // 2
169
+ xmax, ymax = x + w // 2, y + h // 2
170
+ polygon = (xmin, ymin), (xmax, ymin), (xmax, ymax), (xmin, ymax), (xmin, ymin)
171
+ return TableCell(polygon, text_x=x - w // 4, text_y=y, text=value)
172
 
173
 
174
  def all_example_images() -> list[str]:
 
200
  first, last, book = io.read_book(book_path)
201
  book._name = name
202
  book.size_cell = [1.0, 1.0, 1.0, 1.0]
203
+ return (
204
+ [book.read_image(pg) for pg in range(first_page, last_page)],
205
+ book,
206
+ book_path,
207
+ station_tf.read_text(),
208
+ )
209
 
210
 
211
  def overwrite_table_format_file(book: io.Book, book_path, table_fmt: str):
 
218
 
219
 
220
  with gr.Blocks() as submit:
 
221
  gr.Markdown(
222
+ "🛈 Select or upload the image you want to transcribe. You can upload up to five images at a time."
223
  )
224
 
225
  batch_book_state = gr.State()
 
230
  with gr.Row(equal_height=True):
231
  with gr.Column(scale=5):
232
  batch_image_gallery = gr.Gallery(
233
+ # file_types=[".pdf", ".zarr.zip"],
234
+ label="Book to digitize (should be a .pdf or .zarr.zip file)",
235
  interactive=True,
236
  object_fit="scale-down",
237
+ scale=1.0,
238
  )
239
 
240
  with gr.Column(scale=2):
 
248
  object_fit="scale-down",
249
  min_width=250,
250
  )
251
+ upload_button = gr.UploadButton(min_width=200)
252
 
253
  with Modal(visible=False) as edit_table_fmt_modal:
254
  with gr.Column():
 
257
  "Write a custom table format, overriding the default one. "
258
  "Click on the **Save** button when you are done."
259
  )
260
+ save_tf_button = gr.Button(
261
+ "Save", variant="primary", scale=0, min_width=200
262
+ )
263
  gr.HTML(
264
  (
265
  "<a href='https://dawsonia.readthedocs.io/en/latest/user_guide/misc.html#table-formats' target='_blank'>"
 
271
  )
272
  table_fmt_config_override = gr.Code("", language="python")
273
 
274
+ with gr.Row():
275
+ prob_thresh = gr.Slider(
276
+ minimum=0.0,
277
+ maximum=1.0,
278
+ value=0.75,
279
+ step=0.05,
280
+ label="Prediction probability threshold",
281
+ )
282
+
283
  with gr.Row():
284
  run_button = gr.Button("Digitize", variant="primary", scale=0, min_width=200)
285
  edit_table_fmt_button = gr.Button(
 
291
  examples.select(
292
  get_selected_example_image,
293
  (first_page, last_page),
294
+ (
295
+ batch_image_gallery,
296
+ batch_book_state,
297
+ batch_book_path_state,
298
+ table_fmt_config_override,
299
+ ),
300
  trigger_mode="always_last",
301
  )
302
 
 
305
  outputs=[batch_image_gallery],
306
  )
307
  def validate_images(images):
308
+ print(images)
309
  if len(images) > MAX_IMAGES:
310
  gr.Warning(f"Maximum images you can upload is set to: {MAX_IMAGES}")
311
  return gr.update(value=None)
312
+
313
+ gr.Warning(
314
+ "Digitizing uploaded images is not implemented yet! Work in progress!"
315
+ )
316
  raise NotImplementedError("WIP")
317
  return images
318
+
319
  run_button.click(
320
  fn=run_dawsonia,
321
+ inputs=(
322
+ table_fmt_config_override,
323
+ first_page,
324
+ last_page,
325
+ prob_thresh,
326
+ batch_book_state,
327
+ batch_image_gallery,
328
+ ),
329
  outputs=(collection_submit_state, batch_image_gallery),
330
  )
331
 
332
  ## Table formats modal dialog box
333
  edit_table_fmt_button.click(lambda: Modal(visible=True), None, edit_table_fmt_modal)
334
  save_tf_button.click(
335
+ overwrite_table_format_file,
336
+ (batch_book_state, batch_book_path_state, table_fmt_config_override),
337
+ (batch_book_state,),
338
+ )
app/tabs/visualizer.py CHANGED
@@ -26,11 +26,13 @@ class Page(NamedTuple):
26
 
27
  Collection: TypeAlias = list[Page]
28
 
 
29
  def render_image(collection: Collection, current_page_index: int) -> str:
30
  return _IMAGE_TEMPLATE.render(
31
  page=collection[current_page_index],
32
  )
33
 
 
34
  def toggle_navigation_button(collection: Collection):
35
  visible = len(collection) > 1
36
  return gr.update(visible=visible)
@@ -62,9 +64,8 @@ def update_image_caption(collection: Collection, current_page_index):
62
 
63
 
64
  with gr.Blocks() as visualizer:
65
- gr.Markdown("# Result")
66
  gr.Markdown(
67
- "The image to the below shows where Dawsonia found text in the image."
68
  )
69
 
70
  with gr.Row():
@@ -86,7 +87,7 @@ with gr.Blocks() as visualizer:
86
  "← Previous", visible=False, interactive=False, scale=0
87
  )
88
  right = gr.Button("Next →", visible=False, scale=0)
89
-
90
  collection = gr.State()
91
  current_page_index = gr.State(0)
92
 
@@ -95,7 +96,7 @@ with gr.Blocks() as visualizer:
95
  right.click(
96
  right_button_click, [collection, current_page_index], current_page_index
97
  )
98
-
99
  # Updates on collection change:
100
  # - update the view
101
  # - reset the page index (always start on page 0)
 
26
 
27
  Collection: TypeAlias = list[Page]
28
 
29
+
30
  def render_image(collection: Collection, current_page_index: int) -> str:
31
  return _IMAGE_TEMPLATE.render(
32
  page=collection[current_page_index],
33
  )
34
 
35
+
36
  def toggle_navigation_button(collection: Collection):
37
  visible = len(collection) > 1
38
  return gr.update(visible=visible)
 
64
 
65
 
66
  with gr.Blocks() as visualizer:
 
67
  gr.Markdown(
68
+ "🛈 The image to the below shows where Dawsonia found text in the image."
69
  )
70
 
71
  with gr.Row():
 
87
  "← Previous", visible=False, interactive=False, scale=0
88
  )
89
  right = gr.Button("Next →", visible=False, scale=0)
90
+
91
  collection = gr.State()
92
  current_page_index = gr.State(0)
93
 
 
96
  right.click(
97
  right_button_click, [collection, current_page_index], current_page_index
98
  )
99
+
100
  # Updates on collection change:
101
  # - update the view
102
  # - reset the page index (always start on page 0)