arkindex-base-worker 0.3.7rc5__py3-none-any.whl → 0.5.0a1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (60) hide show
  1. {arkindex_base_worker-0.3.7rc5.dist-info → arkindex_base_worker-0.5.0a1.dist-info}/METADATA +18 -19
  2. arkindex_base_worker-0.5.0a1.dist-info/RECORD +61 -0
  3. {arkindex_base_worker-0.3.7rc5.dist-info → arkindex_base_worker-0.5.0a1.dist-info}/WHEEL +1 -1
  4. {arkindex_base_worker-0.3.7rc5.dist-info → arkindex_base_worker-0.5.0a1.dist-info}/top_level.txt +2 -0
  5. arkindex_worker/cache.py +1 -1
  6. arkindex_worker/image.py +167 -2
  7. arkindex_worker/models.py +18 -0
  8. arkindex_worker/utils.py +98 -4
  9. arkindex_worker/worker/__init__.py +117 -218
  10. arkindex_worker/worker/base.py +39 -46
  11. arkindex_worker/worker/classification.py +34 -18
  12. arkindex_worker/worker/corpus.py +86 -0
  13. arkindex_worker/worker/dataset.py +89 -26
  14. arkindex_worker/worker/element.py +352 -91
  15. arkindex_worker/worker/entity.py +13 -11
  16. arkindex_worker/worker/image.py +21 -0
  17. arkindex_worker/worker/metadata.py +26 -16
  18. arkindex_worker/worker/process.py +92 -0
  19. arkindex_worker/worker/task.py +5 -4
  20. arkindex_worker/worker/training.py +25 -10
  21. arkindex_worker/worker/transcription.py +89 -68
  22. arkindex_worker/worker/version.py +3 -1
  23. hooks/pre_gen_project.py +3 -0
  24. tests/__init__.py +8 -0
  25. tests/conftest.py +47 -58
  26. tests/test_base_worker.py +212 -12
  27. tests/test_dataset_worker.py +294 -437
  28. tests/test_elements_worker/{test_classifications.py → test_classification.py} +216 -100
  29. tests/test_elements_worker/test_cli.py +3 -11
  30. tests/test_elements_worker/test_corpus.py +168 -0
  31. tests/test_elements_worker/test_dataset.py +106 -157
  32. tests/test_elements_worker/test_element.py +427 -0
  33. tests/test_elements_worker/test_element_create_multiple.py +715 -0
  34. tests/test_elements_worker/test_element_create_single.py +528 -0
  35. tests/test_elements_worker/test_element_list_children.py +969 -0
  36. tests/test_elements_worker/test_element_list_parents.py +530 -0
  37. tests/test_elements_worker/{test_entities.py → test_entity_create.py} +37 -195
  38. tests/test_elements_worker/test_entity_list_and_check.py +160 -0
  39. tests/test_elements_worker/test_image.py +66 -0
  40. tests/test_elements_worker/test_metadata.py +252 -161
  41. tests/test_elements_worker/test_process.py +89 -0
  42. tests/test_elements_worker/test_task.py +8 -18
  43. tests/test_elements_worker/test_training.py +17 -8
  44. tests/test_elements_worker/test_transcription_create.py +873 -0
  45. tests/test_elements_worker/test_transcription_create_with_elements.py +951 -0
  46. tests/test_elements_worker/test_transcription_list.py +450 -0
  47. tests/test_elements_worker/test_version.py +60 -0
  48. tests/test_elements_worker/test_worker.py +578 -293
  49. tests/test_image.py +542 -209
  50. tests/test_merge.py +1 -2
  51. tests/test_utils.py +89 -4
  52. worker-demo/tests/__init__.py +0 -0
  53. worker-demo/tests/conftest.py +32 -0
  54. worker-demo/tests/test_worker.py +12 -0
  55. worker-demo/worker_demo/__init__.py +6 -0
  56. worker-demo/worker_demo/worker.py +19 -0
  57. arkindex_base_worker-0.3.7rc5.dist-info/RECORD +0 -41
  58. tests/test_elements_worker/test_elements.py +0 -2713
  59. tests/test_elements_worker/test_transcriptions.py +0 -2119
  60. {arkindex_base_worker-0.3.7rc5.dist-info → arkindex_base_worker-0.5.0a1.dist-info}/LICENSE +0 -0
@@ -0,0 +1,951 @@
1
+ import json
2
+ import re
3
+ from uuid import UUID
4
+
5
+ import pytest
6
+ from playhouse.shortcuts import model_to_dict
7
+
8
+ from arkindex.exceptions import ErrorResponse
9
+ from arkindex_worker.cache import CachedElement, CachedTranscription
10
+ from arkindex_worker.models import Element
11
+ from arkindex_worker.utils import DEFAULT_BATCH_SIZE
12
+ from arkindex_worker.worker.transcription import TextOrientation
13
+
14
+ from . import BASE_API_CALLS
15
+
16
+ TRANSCRIPTIONS_SAMPLE = [
17
+ {
18
+ "polygon": [[100, 150], [700, 150], [700, 200], [100, 200]],
19
+ "confidence": 0.5,
20
+ "text": "The",
21
+ },
22
+ {
23
+ "polygon": [[0, 0], [2000, 0], [2000, 3000], [0, 3000]],
24
+ "confidence": 0.75,
25
+ "text": "first",
26
+ "element_confidence": 0.75,
27
+ },
28
+ {
29
+ "polygon": [[1000, 300], [1200, 300], [1200, 500], [1000, 500]],
30
+ "confidence": 0.9,
31
+ "text": "line",
32
+ },
33
+ ]
34
+
35
+
36
+ def test_create_element_transcriptions_wrong_element(mock_elements_worker):
37
+ with pytest.raises(
38
+ AssertionError,
39
+ match="element shouldn't be null and should be an Element or CachedElement",
40
+ ):
41
+ mock_elements_worker.create_element_transcriptions(
42
+ element=None,
43
+ sub_element_type="page",
44
+ transcriptions=TRANSCRIPTIONS_SAMPLE,
45
+ )
46
+
47
+ with pytest.raises(
48
+ AssertionError,
49
+ match="element shouldn't be null and should be an Element or CachedElement",
50
+ ):
51
+ mock_elements_worker.create_element_transcriptions(
52
+ element="not element type",
53
+ sub_element_type="page",
54
+ transcriptions=TRANSCRIPTIONS_SAMPLE,
55
+ )
56
+
57
+
58
+ def test_create_element_transcriptions_wrong_sub_element_type(mock_elements_worker):
59
+ elt = Element({"zone": None})
60
+
61
+ with pytest.raises(
62
+ AssertionError,
63
+ match="sub_element_type shouldn't be null and should be of type str",
64
+ ):
65
+ mock_elements_worker.create_element_transcriptions(
66
+ element=elt,
67
+ sub_element_type=None,
68
+ transcriptions=TRANSCRIPTIONS_SAMPLE,
69
+ )
70
+
71
+ with pytest.raises(
72
+ AssertionError,
73
+ match="sub_element_type shouldn't be null and should be of type str",
74
+ ):
75
+ mock_elements_worker.create_element_transcriptions(
76
+ element=elt,
77
+ sub_element_type=1234,
78
+ transcriptions=TRANSCRIPTIONS_SAMPLE,
79
+ )
80
+
81
+
82
+ def test_create_element_transcriptions_wrong_transcriptions(mock_elements_worker):
83
+ elt = Element({"zone": None})
84
+
85
+ with pytest.raises(
86
+ AssertionError,
87
+ match="transcriptions shouldn't be null and should be of type list",
88
+ ):
89
+ mock_elements_worker.create_element_transcriptions(
90
+ element=elt,
91
+ sub_element_type="page",
92
+ transcriptions=None,
93
+ )
94
+
95
+ with pytest.raises(
96
+ AssertionError,
97
+ match="transcriptions shouldn't be null and should be of type list",
98
+ ):
99
+ mock_elements_worker.create_element_transcriptions(
100
+ element=elt,
101
+ sub_element_type="page",
102
+ transcriptions=1234,
103
+ )
104
+
105
+ with pytest.raises(
106
+ AssertionError,
107
+ match="Transcription at index 1 in transcriptions: text shouldn't be null and should be of type str",
108
+ ):
109
+ mock_elements_worker.create_element_transcriptions(
110
+ element=elt,
111
+ sub_element_type="page",
112
+ transcriptions=[
113
+ {
114
+ "polygon": [[0, 0], [2000, 0], [2000, 3000], [0, 3000]],
115
+ "confidence": 0.75,
116
+ "text": "The",
117
+ },
118
+ {
119
+ "polygon": [[100, 150], [700, 150], [700, 200], [100, 200]],
120
+ "confidence": 0.5,
121
+ },
122
+ ],
123
+ )
124
+
125
+ with pytest.raises(
126
+ AssertionError,
127
+ match="Transcription at index 1 in transcriptions: text shouldn't be null and should be of type str",
128
+ ):
129
+ mock_elements_worker.create_element_transcriptions(
130
+ element=elt,
131
+ sub_element_type="page",
132
+ transcriptions=[
133
+ {
134
+ "polygon": [[0, 0], [2000, 0], [2000, 3000], [0, 3000]],
135
+ "confidence": 0.75,
136
+ "text": "The",
137
+ },
138
+ {
139
+ "polygon": [[100, 150], [700, 150], [700, 200], [100, 200]],
140
+ "confidence": 0.5,
141
+ "text": None,
142
+ },
143
+ ],
144
+ )
145
+
146
+ with pytest.raises(
147
+ AssertionError,
148
+ match="Transcription at index 1 in transcriptions: text shouldn't be null and should be of type str",
149
+ ):
150
+ mock_elements_worker.create_element_transcriptions(
151
+ element=elt,
152
+ sub_element_type="page",
153
+ transcriptions=[
154
+ {
155
+ "polygon": [[0, 0], [2000, 0], [2000, 3000], [0, 3000]],
156
+ "confidence": 0.75,
157
+ "text": "The",
158
+ },
159
+ {
160
+ "polygon": [[100, 150], [700, 150], [700, 200], [100, 200]],
161
+ "confidence": 0.5,
162
+ "text": 1234,
163
+ },
164
+ ],
165
+ )
166
+
167
+ with pytest.raises(
168
+ AssertionError,
169
+ match=re.escape(
170
+ "Transcription at index 1 in transcriptions: confidence shouldn't be null and should be a float in [0..1] range"
171
+ ),
172
+ ):
173
+ mock_elements_worker.create_element_transcriptions(
174
+ element=elt,
175
+ sub_element_type="page",
176
+ transcriptions=[
177
+ {
178
+ "polygon": [[0, 0], [2000, 0], [2000, 3000], [0, 3000]],
179
+ "confidence": 0.75,
180
+ "text": "The",
181
+ },
182
+ {
183
+ "polygon": [[100, 150], [700, 150], [700, 200], [100, 200]],
184
+ "text": "word",
185
+ },
186
+ ],
187
+ )
188
+
189
+ with pytest.raises(
190
+ AssertionError,
191
+ match=re.escape(
192
+ "Transcription at index 1 in transcriptions: confidence shouldn't be null and should be a float in [0..1] range"
193
+ ),
194
+ ):
195
+ mock_elements_worker.create_element_transcriptions(
196
+ element=elt,
197
+ sub_element_type="page",
198
+ transcriptions=[
199
+ {
200
+ "polygon": [[0, 0], [2000, 0], [2000, 3000], [0, 3000]],
201
+ "confidence": 0.75,
202
+ "text": "The",
203
+ },
204
+ {
205
+ "polygon": [[100, 150], [700, 150], [700, 200], [100, 200]],
206
+ "confidence": None,
207
+ "text": "word",
208
+ },
209
+ ],
210
+ )
211
+
212
+ with pytest.raises(
213
+ AssertionError,
214
+ match=re.escape(
215
+ "Transcription at index 1 in transcriptions: confidence shouldn't be null and should be a float in [0..1] range"
216
+ ),
217
+ ):
218
+ mock_elements_worker.create_element_transcriptions(
219
+ element=elt,
220
+ sub_element_type="page",
221
+ transcriptions=[
222
+ {
223
+ "polygon": [[0, 0], [2000, 0], [2000, 3000], [0, 3000]],
224
+ "confidence": 0.75,
225
+ "text": "The",
226
+ },
227
+ {
228
+ "polygon": [[100, 150], [700, 150], [700, 200], [100, 200]],
229
+ "confidence": "a wrong confidence",
230
+ "text": "word",
231
+ },
232
+ ],
233
+ )
234
+
235
+ with pytest.raises(
236
+ AssertionError,
237
+ match=re.escape(
238
+ "Transcription at index 1 in transcriptions: confidence shouldn't be null and should be a float in [0..1] range"
239
+ ),
240
+ ):
241
+ mock_elements_worker.create_element_transcriptions(
242
+ element=elt,
243
+ sub_element_type="page",
244
+ transcriptions=[
245
+ {
246
+ "polygon": [[0, 0], [2000, 0], [2000, 3000], [0, 3000]],
247
+ "confidence": 0.75,
248
+ "text": "The",
249
+ },
250
+ {
251
+ "polygon": [[100, 150], [700, 150], [700, 200], [100, 200]],
252
+ "confidence": 0,
253
+ "text": "word",
254
+ },
255
+ ],
256
+ )
257
+
258
+ with pytest.raises(
259
+ AssertionError,
260
+ match=re.escape(
261
+ "Transcription at index 1 in transcriptions: confidence shouldn't be null and should be a float in [0..1] range"
262
+ ),
263
+ ):
264
+ mock_elements_worker.create_element_transcriptions(
265
+ element=elt,
266
+ sub_element_type="page",
267
+ transcriptions=[
268
+ {
269
+ "polygon": [[0, 0], [2000, 0], [2000, 3000], [0, 3000]],
270
+ "confidence": 0.75,
271
+ "text": "The",
272
+ },
273
+ {
274
+ "polygon": [[100, 150], [700, 150], [700, 200], [100, 200]],
275
+ "confidence": 2.00,
276
+ "text": "word",
277
+ },
278
+ ],
279
+ )
280
+
281
+ with pytest.raises(
282
+ AssertionError,
283
+ match="Transcription at index 1 in transcriptions: polygon shouldn't be null and should be of type list",
284
+ ):
285
+ mock_elements_worker.create_element_transcriptions(
286
+ element=elt,
287
+ sub_element_type="page",
288
+ transcriptions=[
289
+ {
290
+ "polygon": [[0, 0], [2000, 0], [2000, 3000], [0, 3000]],
291
+ "confidence": 0.75,
292
+ "text": "The",
293
+ },
294
+ {"confidence": 0.5, "text": "word"},
295
+ ],
296
+ )
297
+
298
+ with pytest.raises(
299
+ AssertionError,
300
+ match="Transcription at index 1 in transcriptions: polygon shouldn't be null and should be of type list",
301
+ ):
302
+ mock_elements_worker.create_element_transcriptions(
303
+ element=elt,
304
+ sub_element_type="page",
305
+ transcriptions=[
306
+ {
307
+ "polygon": [[0, 0], [2000, 0], [2000, 3000], [0, 3000]],
308
+ "confidence": 0.75,
309
+ "text": "The",
310
+ },
311
+ {"polygon": None, "confidence": 0.5, "text": "word"},
312
+ ],
313
+ )
314
+
315
+ with pytest.raises(
316
+ AssertionError,
317
+ match="Transcription at index 1 in transcriptions: polygon shouldn't be null and should be of type list",
318
+ ):
319
+ mock_elements_worker.create_element_transcriptions(
320
+ element=elt,
321
+ sub_element_type="page",
322
+ transcriptions=[
323
+ {
324
+ "polygon": [[0, 0], [2000, 0], [2000, 3000], [0, 3000]],
325
+ "confidence": 0.75,
326
+ "text": "The",
327
+ },
328
+ {"polygon": "not a polygon", "confidence": 0.5, "text": "word"},
329
+ ],
330
+ )
331
+
332
+ with pytest.raises(
333
+ AssertionError,
334
+ match="Transcription at index 1 in transcriptions: polygon should have at least three points",
335
+ ):
336
+ mock_elements_worker.create_element_transcriptions(
337
+ element=elt,
338
+ sub_element_type="page",
339
+ transcriptions=[
340
+ {
341
+ "polygon": [[0, 0], [2000, 0], [2000, 3000], [0, 3000]],
342
+ "confidence": 0.75,
343
+ "text": "The",
344
+ },
345
+ {"polygon": [[1, 1], [2, 2]], "confidence": 0.5, "text": "word"},
346
+ ],
347
+ )
348
+ with pytest.raises(
349
+ AssertionError,
350
+ match="Transcription at index 1 in transcriptions: polygon points should be lists of two items",
351
+ ):
352
+ mock_elements_worker.create_element_transcriptions(
353
+ element=elt,
354
+ sub_element_type="page",
355
+ transcriptions=[
356
+ {
357
+ "polygon": [[0, 0], [2000, 0], [2000, 3000], [0, 3000]],
358
+ "confidence": 0.75,
359
+ "text": "The",
360
+ },
361
+ {
362
+ "polygon": [[1, 1, 1], [2, 2, 1], [2, 1, 1], [1, 2, 1]],
363
+ "confidence": 0.5,
364
+ "text": "word",
365
+ },
366
+ ],
367
+ )
368
+
369
+ with pytest.raises(
370
+ AssertionError,
371
+ match="Transcription at index 1 in transcriptions: polygon points should be lists of two items",
372
+ ):
373
+ mock_elements_worker.create_element_transcriptions(
374
+ element=elt,
375
+ sub_element_type="page",
376
+ transcriptions=[
377
+ {
378
+ "polygon": [[0, 0], [2000, 0], [2000, 3000], [0, 3000]],
379
+ "confidence": 0.75,
380
+ "text": "The",
381
+ },
382
+ {"polygon": [[1], [2], [2], [1]], "confidence": 0.5, "text": "word"},
383
+ ],
384
+ )
385
+
386
+ with pytest.raises(
387
+ AssertionError,
388
+ match="Transcription at index 1 in transcriptions: polygon points should be lists of two numbers",
389
+ ):
390
+ mock_elements_worker.create_element_transcriptions(
391
+ element=elt,
392
+ sub_element_type="page",
393
+ transcriptions=[
394
+ {
395
+ "polygon": [[0, 0], [2000, 0], [2000, 3000], [0, 3000]],
396
+ "confidence": 0.75,
397
+ "text": "The",
398
+ },
399
+ {
400
+ "polygon": [["not a coord", 1], [2, 2], [2, 1], [1, 2]],
401
+ "confidence": 0.5,
402
+ "text": "word",
403
+ },
404
+ ],
405
+ )
406
+
407
+ with pytest.raises(
408
+ AssertionError,
409
+ match="Transcription at index 1 in transcriptions: orientation shouldn't be null and should be of type TextOrientation",
410
+ ):
411
+ mock_elements_worker.create_element_transcriptions(
412
+ element=elt,
413
+ sub_element_type="page",
414
+ transcriptions=[
415
+ {
416
+ "polygon": [[0, 0], [2000, 0], [2000, 3000], [0, 3000]],
417
+ "confidence": 0.75,
418
+ "text": "The",
419
+ },
420
+ {
421
+ "polygon": [[100, 150], [700, 150], [700, 200], [100, 200]],
422
+ "confidence": 0.35,
423
+ "text": "word",
424
+ "orientation": "uptown",
425
+ },
426
+ ],
427
+ )
428
+
429
+ with pytest.raises(
430
+ AssertionError,
431
+ match=re.escape(
432
+ "Transcription at index 1 in transcriptions: element_confidence should be either null or a float in [0..1] range"
433
+ ),
434
+ ):
435
+ mock_elements_worker.create_element_transcriptions(
436
+ element=elt,
437
+ sub_element_type="page",
438
+ transcriptions=[
439
+ {
440
+ "polygon": [[0, 0], [2000, 0], [2000, 3000], [0, 3000]],
441
+ "confidence": 0.75,
442
+ "text": "The",
443
+ },
444
+ {
445
+ "polygon": [[100, 150], [700, 150], [700, 200], [100, 200]],
446
+ "confidence": 0.75,
447
+ "text": "word",
448
+ "element_confidence": "not a confidence",
449
+ },
450
+ ],
451
+ )
452
+
453
+
454
+ def test_create_element_transcriptions_api_error(responses, mock_elements_worker):
455
+ elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
456
+ responses.add(
457
+ responses.POST,
458
+ f"http://testserver/api/v1/element/{elt.id}/transcriptions/bulk/",
459
+ status=418,
460
+ )
461
+
462
+ with pytest.raises(ErrorResponse):
463
+ mock_elements_worker.create_element_transcriptions(
464
+ element=elt,
465
+ sub_element_type="page",
466
+ transcriptions=TRANSCRIPTIONS_SAMPLE,
467
+ )
468
+
469
+ assert len(responses.calls) == len(BASE_API_CALLS) + 1
470
+ assert [
471
+ (call.request.method, call.request.url) for call in responses.calls
472
+ ] == BASE_API_CALLS + [
473
+ ("POST", f"http://testserver/api/v1/element/{elt.id}/transcriptions/bulk/")
474
+ ]
475
+
476
+
477
+ @pytest.mark.parametrize("batch_size", [DEFAULT_BATCH_SIZE, 2])
478
+ def test_create_element_transcriptions(batch_size, responses, mock_elements_worker):
479
+ elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
480
+
481
+ if batch_size > 2:
482
+ responses.add(
483
+ responses.POST,
484
+ f"http://testserver/api/v1/element/{elt.id}/transcriptions/bulk/",
485
+ status=200,
486
+ json=[
487
+ {
488
+ "id": "56785678-5678-5678-5678-567856785678",
489
+ "element_id": "11111111-1111-1111-1111-111111111111",
490
+ "created": True,
491
+ },
492
+ {
493
+ "id": "67896789-6789-6789-6789-678967896789",
494
+ "element_id": "22222222-2222-2222-2222-222222222222",
495
+ "created": False,
496
+ },
497
+ {
498
+ "id": "78907890-7890-7890-7890-789078907890",
499
+ "element_id": "11111111-1111-1111-1111-111111111111",
500
+ "created": True,
501
+ },
502
+ ],
503
+ )
504
+ else:
505
+ for transcriptions in [
506
+ [
507
+ ("56785678-5678-5678-5678-567856785678", True),
508
+ ("67896789-6789-6789-6789-678967896789", False),
509
+ ],
510
+ [("78907890-7890-7890-7890-789078907890", True)],
511
+ ]:
512
+ responses.add(
513
+ responses.POST,
514
+ f"http://testserver/api/v1/element/{elt.id}/transcriptions/bulk/",
515
+ status=200,
516
+ json=[
517
+ {
518
+ "id": tr_id,
519
+ "element_id": "11111111-1111-1111-1111-111111111111"
520
+ if created
521
+ else "22222222-2222-2222-2222-222222222222",
522
+ "created": created,
523
+ }
524
+ for tr_id, created in transcriptions
525
+ ],
526
+ )
527
+
528
+ annotations = mock_elements_worker.create_element_transcriptions(
529
+ element=elt,
530
+ sub_element_type="page",
531
+ transcriptions=TRANSCRIPTIONS_SAMPLE,
532
+ batch_size=batch_size,
533
+ )
534
+
535
+ bulk_api_calls = [
536
+ (
537
+ "POST",
538
+ f"http://testserver/api/v1/element/{elt.id}/transcriptions/bulk/",
539
+ )
540
+ ]
541
+ if batch_size != DEFAULT_BATCH_SIZE:
542
+ bulk_api_calls.append(
543
+ (
544
+ "POST",
545
+ f"http://testserver/api/v1/element/{elt.id}/transcriptions/bulk/",
546
+ )
547
+ )
548
+
549
+ assert len(responses.calls) == len(BASE_API_CALLS) + len(bulk_api_calls)
550
+ assert [
551
+ (call.request.method, call.request.url) for call in responses.calls
552
+ ] == BASE_API_CALLS + bulk_api_calls
553
+
554
+ first_tr = {
555
+ **TRANSCRIPTIONS_SAMPLE[0],
556
+ "orientation": TextOrientation.HorizontalLeftToRight.value,
557
+ }
558
+ second_tr = {
559
+ **TRANSCRIPTIONS_SAMPLE[1],
560
+ "orientation": TextOrientation.HorizontalLeftToRight.value,
561
+ }
562
+ third_tr = {
563
+ **TRANSCRIPTIONS_SAMPLE[2],
564
+ "orientation": TextOrientation.HorizontalLeftToRight.value,
565
+ }
566
+ empty_payload = {
567
+ "element_type": "page",
568
+ "worker_run_id": "56785678-5678-5678-5678-567856785678",
569
+ "transcriptions": [],
570
+ "return_elements": True,
571
+ }
572
+
573
+ bodies = []
574
+ first_call_idx = None
575
+ if batch_size > 2:
576
+ first_call_idx = -1
577
+ bodies.append(
578
+ {**empty_payload, "transcriptions": [first_tr, second_tr, third_tr]}
579
+ )
580
+ else:
581
+ first_call_idx = -2
582
+ bodies.append({**empty_payload, "transcriptions": [first_tr, second_tr]})
583
+ bodies.append({**empty_payload, "transcriptions": [third_tr]})
584
+
585
+ assert [
586
+ json.loads(bulk_call.request.body)
587
+ for bulk_call in responses.calls[first_call_idx:]
588
+ ] == bodies
589
+
590
+ assert annotations == [
591
+ {
592
+ "id": "56785678-5678-5678-5678-567856785678",
593
+ "element_id": "11111111-1111-1111-1111-111111111111",
594
+ "created": True,
595
+ },
596
+ {
597
+ "id": "67896789-6789-6789-6789-678967896789",
598
+ "element_id": "22222222-2222-2222-2222-222222222222",
599
+ "created": False,
600
+ },
601
+ {
602
+ "id": "78907890-7890-7890-7890-789078907890",
603
+ "element_id": "11111111-1111-1111-1111-111111111111",
604
+ "created": True,
605
+ },
606
+ ]
607
+
608
+
609
+ @pytest.mark.parametrize("batch_size", [DEFAULT_BATCH_SIZE, 2])
610
+ def test_create_element_transcriptions_with_cache(
611
+ batch_size, responses, mock_elements_worker_with_cache
612
+ ):
613
+ elt = CachedElement(id="12341234-1234-1234-1234-123412341234", type="thing")
614
+
615
+ if batch_size > 2:
616
+ responses.add(
617
+ responses.POST,
618
+ f"http://testserver/api/v1/element/{elt.id}/transcriptions/bulk/",
619
+ status=200,
620
+ json=[
621
+ {
622
+ "id": "56785678-5678-5678-5678-567856785678",
623
+ "element_id": "11111111-1111-1111-1111-111111111111",
624
+ "created": True,
625
+ },
626
+ {
627
+ "id": "67896789-6789-6789-6789-678967896789",
628
+ "element_id": "22222222-2222-2222-2222-222222222222",
629
+ "created": False,
630
+ },
631
+ {
632
+ "id": "78907890-7890-7890-7890-789078907890",
633
+ "element_id": "11111111-1111-1111-1111-111111111111",
634
+ "created": True,
635
+ },
636
+ ],
637
+ )
638
+ else:
639
+ for transcriptions in [
640
+ [
641
+ ("56785678-5678-5678-5678-567856785678", True),
642
+ ("67896789-6789-6789-6789-678967896789", False),
643
+ ],
644
+ [("78907890-7890-7890-7890-789078907890", True)],
645
+ ]:
646
+ responses.add(
647
+ responses.POST,
648
+ f"http://testserver/api/v1/element/{elt.id}/transcriptions/bulk/",
649
+ status=200,
650
+ json=[
651
+ {
652
+ "id": tr_id,
653
+ "element_id": "11111111-1111-1111-1111-111111111111"
654
+ if created
655
+ else "22222222-2222-2222-2222-222222222222",
656
+ "created": created,
657
+ }
658
+ for tr_id, created in transcriptions
659
+ ],
660
+ )
661
+
662
+ annotations = mock_elements_worker_with_cache.create_element_transcriptions(
663
+ element=elt,
664
+ sub_element_type="page",
665
+ transcriptions=TRANSCRIPTIONS_SAMPLE,
666
+ batch_size=batch_size,
667
+ )
668
+
669
+ bulk_api_calls = [
670
+ (
671
+ "POST",
672
+ f"http://testserver/api/v1/element/{elt.id}/transcriptions/bulk/",
673
+ )
674
+ ]
675
+ if batch_size != DEFAULT_BATCH_SIZE:
676
+ bulk_api_calls.append(
677
+ (
678
+ "POST",
679
+ f"http://testserver/api/v1/element/{elt.id}/transcriptions/bulk/",
680
+ )
681
+ )
682
+
683
+ assert len(responses.calls) == len(BASE_API_CALLS) + len(bulk_api_calls)
684
+ assert [
685
+ (call.request.method, call.request.url) for call in responses.calls
686
+ ] == BASE_API_CALLS + bulk_api_calls
687
+
688
+ first_tr = {
689
+ **TRANSCRIPTIONS_SAMPLE[0],
690
+ "orientation": TextOrientation.HorizontalLeftToRight.value,
691
+ }
692
+ second_tr = {
693
+ **TRANSCRIPTIONS_SAMPLE[1],
694
+ "orientation": TextOrientation.HorizontalLeftToRight.value,
695
+ "element_confidence": 0.75,
696
+ }
697
+ third_tr = {
698
+ **TRANSCRIPTIONS_SAMPLE[2],
699
+ "orientation": TextOrientation.HorizontalLeftToRight.value,
700
+ }
701
+ empty_payload = {
702
+ "element_type": "page",
703
+ "worker_run_id": "56785678-5678-5678-5678-567856785678",
704
+ "transcriptions": [],
705
+ "return_elements": True,
706
+ }
707
+
708
+ bodies = []
709
+ first_call_idx = None
710
+ if batch_size > 2:
711
+ first_call_idx = -1
712
+ bodies.append(
713
+ {**empty_payload, "transcriptions": [first_tr, second_tr, third_tr]}
714
+ )
715
+ else:
716
+ first_call_idx = -2
717
+ bodies.append({**empty_payload, "transcriptions": [first_tr, second_tr]})
718
+ bodies.append({**empty_payload, "transcriptions": [third_tr]})
719
+
720
+ assert [
721
+ json.loads(bulk_call.request.body)
722
+ for bulk_call in responses.calls[first_call_idx:]
723
+ ] == bodies
724
+
725
+ assert annotations == [
726
+ {
727
+ "id": "56785678-5678-5678-5678-567856785678",
728
+ "element_id": "11111111-1111-1111-1111-111111111111",
729
+ "created": True,
730
+ },
731
+ {
732
+ "id": "67896789-6789-6789-6789-678967896789",
733
+ "element_id": "22222222-2222-2222-2222-222222222222",
734
+ "created": False,
735
+ },
736
+ {
737
+ "id": "78907890-7890-7890-7890-789078907890",
738
+ "element_id": "11111111-1111-1111-1111-111111111111",
739
+ "created": True,
740
+ },
741
+ ]
742
+
743
+ # Check that created transcriptions and elements were properly stored in SQLite cache
744
+ assert list(CachedElement.select()) == [
745
+ CachedElement(
746
+ id=UUID("11111111-1111-1111-1111-111111111111"),
747
+ parent_id=UUID("12341234-1234-1234-1234-123412341234"),
748
+ type="page",
749
+ polygon="[[100, 150], [700, 150], [700, 200], [100, 200]]",
750
+ worker_run_id=UUID("56785678-5678-5678-5678-567856785678"),
751
+ ),
752
+ CachedElement(
753
+ id=UUID("22222222-2222-2222-2222-222222222222"),
754
+ parent_id=UUID("12341234-1234-1234-1234-123412341234"),
755
+ type="page",
756
+ polygon="[[0, 0], [2000, 0], [2000, 3000], [0, 3000]]",
757
+ worker_run_id=UUID("56785678-5678-5678-5678-567856785678"),
758
+ confidence=0.75,
759
+ ),
760
+ ]
761
+ assert list(CachedTranscription.select()) == [
762
+ CachedTranscription(
763
+ id=UUID("56785678-5678-5678-5678-567856785678"),
764
+ element_id=UUID("11111111-1111-1111-1111-111111111111"),
765
+ text="The",
766
+ confidence=0.5,
767
+ orientation=TextOrientation.HorizontalLeftToRight.value,
768
+ worker_run_id=UUID("56785678-5678-5678-5678-567856785678"),
769
+ ),
770
+ CachedTranscription(
771
+ id=UUID("67896789-6789-6789-6789-678967896789"),
772
+ element_id=UUID("22222222-2222-2222-2222-222222222222"),
773
+ text="first",
774
+ confidence=0.75,
775
+ orientation=TextOrientation.HorizontalLeftToRight.value,
776
+ worker_run_id=UUID("56785678-5678-5678-5678-567856785678"),
777
+ ),
778
+ CachedTranscription(
779
+ id=UUID("78907890-7890-7890-7890-789078907890"),
780
+ element_id=UUID("11111111-1111-1111-1111-111111111111"),
781
+ text="line",
782
+ confidence=0.9,
783
+ orientation=TextOrientation.HorizontalLeftToRight.value,
784
+ worker_run_id=UUID("56785678-5678-5678-5678-567856785678"),
785
+ ),
786
+ ]
787
+
788
+
789
+ def test_create_element_transcriptions_orientation_with_cache(
790
+ responses, mock_elements_worker_with_cache
791
+ ):
792
+ elt = CachedElement(id="12341234-1234-1234-1234-123412341234", type="thing")
793
+
794
+ responses.add(
795
+ responses.POST,
796
+ f"http://testserver/api/v1/element/{elt.id}/transcriptions/bulk/",
797
+ status=200,
798
+ json=[
799
+ {
800
+ "id": "56785678-5678-5678-5678-567856785678",
801
+ "element_id": "11111111-1111-1111-1111-111111111111",
802
+ "created": True,
803
+ },
804
+ {
805
+ "id": "67896789-6789-6789-6789-678967896789",
806
+ "element_id": "22222222-2222-2222-2222-222222222222",
807
+ "created": False,
808
+ },
809
+ {
810
+ "id": "78907890-7890-7890-7890-789078907890",
811
+ "element_id": "11111111-1111-1111-1111-111111111111",
812
+ "created": True,
813
+ },
814
+ ],
815
+ )
816
+
817
+ oriented_transcriptions = [
818
+ {
819
+ "polygon": [[100, 150], [700, 150], [700, 200], [100, 200]],
820
+ "confidence": 0.5,
821
+ "text": "Animula vagula blandula",
822
+ },
823
+ {
824
+ "polygon": [[0, 0], [2000, 0], [2000, 3000], [0, 3000]],
825
+ "confidence": 0.75,
826
+ "text": "Hospes comesque corporis",
827
+ "orientation": TextOrientation.VerticalLeftToRight,
828
+ },
829
+ {
830
+ "polygon": [[100, 150], [700, 150], [700, 200], [100, 200]],
831
+ "confidence": 0.9,
832
+ "text": "Quae nunc abibis in loca",
833
+ "orientation": TextOrientation.HorizontalRightToLeft,
834
+ },
835
+ ]
836
+
837
+ annotations = mock_elements_worker_with_cache.create_element_transcriptions(
838
+ element=elt,
839
+ sub_element_type="page",
840
+ transcriptions=oriented_transcriptions,
841
+ )
842
+
843
+ assert json.loads(responses.calls[-1].request.body) == {
844
+ "element_type": "page",
845
+ "worker_run_id": "56785678-5678-5678-5678-567856785678",
846
+ "transcriptions": [
847
+ {
848
+ "polygon": [[100, 150], [700, 150], [700, 200], [100, 200]],
849
+ "confidence": 0.5,
850
+ "text": "Animula vagula blandula",
851
+ "orientation": TextOrientation.HorizontalLeftToRight.value,
852
+ },
853
+ {
854
+ "polygon": [[0, 0], [2000, 0], [2000, 3000], [0, 3000]],
855
+ "confidence": 0.75,
856
+ "text": "Hospes comesque corporis",
857
+ "orientation": TextOrientation.VerticalLeftToRight.value,
858
+ },
859
+ {
860
+ "polygon": [[100, 150], [700, 150], [700, 200], [100, 200]],
861
+ "confidence": 0.9,
862
+ "text": "Quae nunc abibis in loca",
863
+ "orientation": TextOrientation.HorizontalRightToLeft.value,
864
+ },
865
+ ],
866
+ "return_elements": True,
867
+ }
868
+ assert annotations == [
869
+ {
870
+ "id": "56785678-5678-5678-5678-567856785678",
871
+ "element_id": "11111111-1111-1111-1111-111111111111",
872
+ "created": True,
873
+ },
874
+ {
875
+ "id": "67896789-6789-6789-6789-678967896789",
876
+ "element_id": "22222222-2222-2222-2222-222222222222",
877
+ "created": False,
878
+ },
879
+ {
880
+ "id": "78907890-7890-7890-7890-789078907890",
881
+ "element_id": "11111111-1111-1111-1111-111111111111",
882
+ "created": True,
883
+ },
884
+ ]
885
+
886
+ # Check that the text orientation was properly stored in SQLite cache
887
+ assert list(map(model_to_dict, CachedTranscription.select())) == [
888
+ {
889
+ "id": UUID("56785678-5678-5678-5678-567856785678"),
890
+ "element": {
891
+ "id": UUID("11111111-1111-1111-1111-111111111111"),
892
+ "parent_id": UUID(elt.id),
893
+ "type": "page",
894
+ "image": None,
895
+ "polygon": [[100, 150], [700, 150], [700, 200], [100, 200]],
896
+ "rotation_angle": 0,
897
+ "mirrored": False,
898
+ "initial": False,
899
+ "worker_version_id": None,
900
+ "worker_run_id": UUID("56785678-5678-5678-5678-567856785678"),
901
+ "confidence": None,
902
+ },
903
+ "text": "Animula vagula blandula",
904
+ "confidence": 0.5,
905
+ "orientation": TextOrientation.HorizontalLeftToRight.value,
906
+ "worker_version_id": None,
907
+ "worker_run_id": UUID("56785678-5678-5678-5678-567856785678"),
908
+ },
909
+ {
910
+ "id": UUID("67896789-6789-6789-6789-678967896789"),
911
+ "element": {
912
+ "id": UUID("22222222-2222-2222-2222-222222222222"),
913
+ "parent_id": UUID(elt.id),
914
+ "type": "page",
915
+ "image": None,
916
+ "polygon": [[0, 0], [2000, 0], [2000, 3000], [0, 3000]],
917
+ "rotation_angle": 0,
918
+ "mirrored": False,
919
+ "initial": False,
920
+ "worker_version_id": None,
921
+ "worker_run_id": UUID("56785678-5678-5678-5678-567856785678"),
922
+ "confidence": None,
923
+ },
924
+ "text": "Hospes comesque corporis",
925
+ "confidence": 0.75,
926
+ "orientation": TextOrientation.VerticalLeftToRight.value,
927
+ "worker_version_id": None,
928
+ "worker_run_id": UUID("56785678-5678-5678-5678-567856785678"),
929
+ },
930
+ {
931
+ "id": UUID("78907890-7890-7890-7890-789078907890"),
932
+ "element": {
933
+ "id": UUID("11111111-1111-1111-1111-111111111111"),
934
+ "parent_id": UUID(elt.id),
935
+ "type": "page",
936
+ "image": None,
937
+ "polygon": [[100, 150], [700, 150], [700, 200], [100, 200]],
938
+ "rotation_angle": 0,
939
+ "mirrored": False,
940
+ "initial": False,
941
+ "worker_version_id": None,
942
+ "worker_run_id": UUID("56785678-5678-5678-5678-567856785678"),
943
+ "confidence": None,
944
+ },
945
+ "text": "Quae nunc abibis in loca",
946
+ "confidence": 0.9,
947
+ "orientation": TextOrientation.HorizontalRightToLeft.value,
948
+ "worker_version_id": None,
949
+ "worker_run_id": UUID("56785678-5678-5678-5678-567856785678"),
950
+ },
951
+ ]