OCR Team
293 lines
16 KiB
Python
293 lines
16 KiB
Python
from io import BytesIO
|
|
import json
|
|
from pathlib import Path
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
from types import SimpleNamespace
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
from PIL import Image
|
|
from starlette.testclient import TestClient
|
|
|
|
from ocr_compare import server
|
|
from ocr_compare import classic_worker
|
|
from ocr_compare.fields import assign_fields
|
|
from ocr_compare.vision_worker import _missing_neighbor_box, _retry_table, extract
|
|
|
|
|
|
def photo(width=800, height=500):
|
|
stream = BytesIO()
|
|
Image.new("RGB", (width, height), "white").save(stream, "JPEG")
|
|
return stream.getvalue()
|
|
|
|
|
|
class LocalComparisonTests(unittest.TestCase):
|
|
def test_adjacent_field_code_needs_an_observed_value(self):
|
|
lines = [
|
|
{"text": "D.3", "box": [10, 100, 35, 120]},
|
|
{"text": "ZXR 400", "box": [31, 96, 150, 130]},
|
|
{"text": "P.3", "box": [10, 200, 36, 220]},
|
|
{"text": "BENZIN", "box": [32, 197, 125, 228]},
|
|
{"text": "R", "box": [10, 290, 30, 310]},
|
|
]
|
|
found = {item["code"]: item for item in assign_fields(lines)}
|
|
self.assertEqual(found["D.3"]["value"], "ZXR 400")
|
|
self.assertEqual(found["P.3"]["value"], "BENZIN")
|
|
self.assertNotIn("R", found)
|
|
self.assertEqual(found["D.3"]["box"], [31, 96, 150, 130])
|
|
self.assertIsNone(found["D.3"]["confidence"])
|
|
|
|
def test_merged_code_is_reviewable_and_value_box_is_not_reused(self):
|
|
lines = [
|
|
{"text": "D3330D", "box": [100, 100, 230, 135]},
|
|
{"text": "A", "box": [100, 200, 120, 230]},
|
|
{"text": "16", "box": [135, 200, 160, 230]},
|
|
{"text": "TESTNUMMER", "box": [170, 200, 320, 230]},
|
|
{"text": "C.1.2 Vorname(n)", "box": [10, 300, 240, 330]},
|
|
]
|
|
found = {item["code"]: item for item in assign_fields(lines)}
|
|
self.assertEqual(found["D.3"]["value"], "330D")
|
|
self.assertEqual(found["D.3"]["method"], "merged_code")
|
|
self.assertEqual(found["16"]["value"], "TESTNUMMER")
|
|
self.assertNotIn("A", found)
|
|
self.assertNotIn("C.1.2", found)
|
|
|
|
def test_dotted_type_codes_are_distinct_from_ambiguous_undotted_codes(self):
|
|
lines = [
|
|
{"text": "2.1", "box": [100, 100, 130, 120]},
|
|
{"text": "0005", "box": [140, 100, 195, 120]},
|
|
{"text": "2.2", "box": [210, 100, 240, 120]},
|
|
{"text": "ABC123", "box": [250, 100, 330, 120]},
|
|
{"text": "21", "box": [100, 200, 130, 220]},
|
|
{"text": "UNSICHER", "box": [140, 200, 230, 220]},
|
|
]
|
|
found = {item["code"]: item for item in assign_fields(lines)}
|
|
self.assertEqual(found["2.1"]["value"], "0005")
|
|
self.assertEqual(found["2.2"]["value"], "ABC123")
|
|
self.assertNotIn("21", found)
|
|
|
|
def test_undotted_type_codes_need_the_observed_b_table_sequence(self):
|
|
lines = [{"text": "table", "cells": ["B", "01.02.2003", "21", "0005", "22", "ABC123"],
|
|
"box": [10, 10, 600, 120]}]
|
|
found = {item["code"]: item for item in assign_fields(lines, allow_adjacent=False)}
|
|
self.assertEqual(found["2.1"]["value"], "0005")
|
|
self.assertEqual(found["2.2"]["value"], "ABC123")
|
|
self.assertEqual(found["2.1"]["method"], "vl_table_context_code")
|
|
self.assertNotIn("21", found)
|
|
|
|
def test_holder_values_below_captions_join_only_touching_rows(self):
|
|
lines = [
|
|
{"text": "C.1.1 Name oder Firmenname", "box": [10, 10, 260, 40]},
|
|
{"text": "MUSTER", "box": [20, 55, 180, 95]},
|
|
{"text": "C.1.2 Vorname(n)", "box": [10, 150, 210, 180]},
|
|
{"text": "ANNA", "box": [20, 195, 130, 235]},
|
|
{"text": "MARIA", "box": [125, 195, 250, 235]},
|
|
{"text": "C.1.3 Anschrift", "box": [10, 280, 170, 310]},
|
|
{"text": "TESTSTRASSE", "box": [20, 325, 240, 365]},
|
|
{"text": "6", "box": [235, 325, 255, 365]},
|
|
{"text": "12345", "box": [20, 370, 100, 410]},
|
|
{"text": "TESTSTADT", "box": [105, 370, 250, 410]},
|
|
{"text": "X Naechste HU", "box": [10, 480, 180, 520]},
|
|
]
|
|
found = {item["code"]: item for item in assign_fields(lines)}
|
|
self.assertEqual(found["C.1.1"]["value"], "MUSTER")
|
|
self.assertEqual(found["C.1.2"]["value"], "ANNA MARIA")
|
|
self.assertEqual(found["C.1.3"]["value"], "TESTSTRASSE 6\n12345 TESTSTADT")
|
|
self.assertEqual(found["C.1.3"]["box"], [20, 325, 255, 410])
|
|
|
|
def test_vision_caption_and_two_lines_in_one_layout_region(self):
|
|
lines = [
|
|
{"text": "C.I.I Name oder Firmenname", "box": [10, 10, 250, 40]},
|
|
{"text": "MUSTER", "box": [20, 55, 130, 95]},
|
|
{"text": "C.1.3 Anschrift", "box": [10, 150, 180, 180]},
|
|
{"text": "TESTSTRASSE 6", "box": [20, 195, 260, 270]},
|
|
{"text": "12345 TESTSTADT", "box": [20, 195, 260, 270]},
|
|
{"text": "X Naechste HU", "box": [10, 300, 180, 350]},
|
|
]
|
|
found = {item["code"]: item for item in assign_fields(lines, allow_adjacent=False)}
|
|
self.assertEqual(found["C.1.1"]["value"], "MUSTER")
|
|
self.assertEqual(found["C.1.3"]["value"], "TESTSTRASSE 6\n12345 TESTSTADT")
|
|
|
|
def test_missing_surname_caption_is_only_a_reviewable_nearby_suggestion(self):
|
|
lines = [
|
|
{"text": "MUSTER", "box": [30, 100, 170, 145]},
|
|
{"text": "C.1.2 Vorname(n)", "box": [20, 210, 220, 245]},
|
|
{"text": "ANNA", "box": [30, 265, 140, 305]},
|
|
]
|
|
found = {item["code"]: item for item in assign_fields(lines)}
|
|
self.assertEqual(found["C.1.1"]["method"], "surname_above_forename")
|
|
self.assertEqual(found["C.1.1"]["value"], "MUSTER")
|
|
self.assertNotIn("C.1.1", {item["code"] for item in assign_fields(lines, allow_adjacent=False)})
|
|
|
|
def test_vision_same_region_prefers_next_line_over_a_later_address(self):
|
|
source = "PaddleOCR-VL-Layoutblock text; keine eigene Zeilenbox"
|
|
lines = [
|
|
{"text": "C.1.2 Vorname(n)", "box": [10, 100, 250, 200], "source": source},
|
|
{"text": "ANNA", "box": [10, 100, 250, 200], "source": source},
|
|
{"text": "TESTSTRASSE", "box": [10, 220, 300, 270], "source": source},
|
|
]
|
|
found = {item["code"]: item for item in assign_fields(lines, allow_adjacent=False)}
|
|
self.assertEqual(found["C.1.2"]["value"], "ANNA")
|
|
self.assertEqual(found["C.1.2"]["method"], "caption_region_next_line")
|
|
|
|
def test_vision_address_does_not_absorb_the_next_layout_region(self):
|
|
source = "PaddleOCR-VL-Layoutblock text; keine eigene Zeilenbox"
|
|
lines = [
|
|
{"text": "C.1.3 Anschrift", "box": [10, 100, 180, 140], "source": source},
|
|
{"text": "TESTSTRASSE 6 12345 TESTSTADT", "box": [10, 160, 310, 240], "source": source},
|
|
{"text": "Nächste HU", "box": [10, 300, 200, 350], "source": source},
|
|
]
|
|
found = {item["code"]: item for item in assign_fields(lines, allow_adjacent=False)}
|
|
self.assertEqual(found["C.1.3"]["value"], "TESTSTRASSE 6 12345 TESTSTADT")
|
|
|
|
def test_large_empty_table_is_retried_but_populated_one_is_not(self):
|
|
block = {"block_label": "table", "block_bbox": [100, 10, 850, 790],
|
|
"block_content": "<table></table>"}
|
|
self.assertTrue(_retry_table(block, 1000, 800))
|
|
self.assertFalse(_retry_table({**block, "block_content": "TEXT " * 30}, 1000, 800))
|
|
self.assertFalse(_retry_table({**block, "block_bbox": [100, 10, 500, 400]}, 1000, 800))
|
|
|
|
def test_missing_neighbor_retry_requires_an_uncovered_column(self):
|
|
right = {"block_label": "table", "block_bbox": [600, 200, 860, 650]}
|
|
box = _missing_neighbor_box(right, [right], 1000, 800)
|
|
self.assertIsNotNone(box)
|
|
self.assertLess(box[0], 600)
|
|
covered = {"block_label": "table", "block_bbox": [290, 200, 600, 650]}
|
|
self.assertIsNone(_missing_neighbor_box(right, [right, covered], 1000, 800))
|
|
|
|
def test_overlapping_vision_tiles_do_not_erase_identical_field_readings(self):
|
|
lines = [
|
|
{"text": "D.1 | TESTMARKE", "cells": ["D.1", "TESTMARKE"], "box": [100, 0, 500, 500]},
|
|
{"text": "D.1 | TESTMARKE", "cells": ["D.1", "TESTMARKE"], "box": [200, 0, 600, 500]},
|
|
]
|
|
fields = assign_fields(lines, allow_adjacent=False)
|
|
self.assertEqual([(item["code"], item["value"]) for item in fields], [("D.1", "TESTMARKE")])
|
|
|
|
def test_wide_vision_timeout_retries_same_image_in_local_tiles(self):
|
|
payload = {"lines": [{"text": "D.1 | TESTMARKE", "cells": ["D.1", "TESTMARKE"],
|
|
"box": [10, 10, 100, 100]}], "regions": [], "retry_used": True}
|
|
timed_out = subprocess.CompletedProcess([], 124, b"")
|
|
tiled = subprocess.CompletedProcess([], 0, json.dumps(payload).encode())
|
|
with patch.dict(server.os.environ, {"OCR_VISION_MODE": "ssh",
|
|
"OCR_VISION_SSH_TARGET": "test@example.invalid",
|
|
"OCR_VISION_REMOTE_DIR": "/tmp/test"}):
|
|
with patch.object(server.subprocess, "run", side_effect=[timed_out, tiled]) as remote:
|
|
result = server._vision(photo(3600, 1800))
|
|
self.assertEqual(result["status"], "ok")
|
|
self.assertEqual(result["fields"][0]["value"], "TESTMARKE")
|
|
self.assertEqual(remote.call_count, 2)
|
|
self.assertIn("--tiles-only", remote.call_args.args[0][-1])
|
|
|
|
def test_corners_allow_perspective_and_reject_crossing(self):
|
|
source = photo(1200, 900)
|
|
upright = {"corners": [[.1, .08], [.9, .13], [.84, .91], [.13, .85]], "rotation": 0}
|
|
output, (width, height) = server._prepare(source, json.dumps(upright))
|
|
self.assertGreater(len(output), 1000)
|
|
self.assertGreater(width, 700)
|
|
self.assertGreater(height, 600)
|
|
crossed = {"corners": [[.1, .1], [.9, .9], [.9, .1], [.1, .9]], "rotation": 0}
|
|
with self.assertRaisesRegex(server.InputError, "invalid_corners"):
|
|
server._prepare(source, json.dumps(crossed))
|
|
|
|
def test_realistic_phone_resolution_is_allowed(self):
|
|
output, size = server._prepare(photo(4096, 3072), None)
|
|
self.assertEqual(size, (4000, 3000))
|
|
self.assertGreater(len(output), 1000)
|
|
|
|
def test_one_failed_engine_does_not_hide_the_other(self):
|
|
classic = {"status": "ok", "model": "test classic", "elapsed_ms": 1,
|
|
"lines": [{"text": "P.3 BENZIN", "box": [1, 1, 150, 30]}],
|
|
"fields": [], "regions": []}
|
|
vision = {"status": "error", "error": "vision_failed", "elapsed_ms": 2}
|
|
with patch.object(server, "_classic", return_value=classic), patch.object(server, "_vision", return_value=vision):
|
|
with TestClient(server.app) as client:
|
|
response = client.post("/api/compare", content=photo(), headers={"content-type": "image/jpeg"})
|
|
self.assertEqual(response.status_code, 200)
|
|
self.assertEqual(response.json()["results"]["classic"], classic)
|
|
self.assertEqual(response.json()["results"]["vision"], vision)
|
|
|
|
def test_fine_text_choice_reaches_classic_engine_only(self):
|
|
classic = {"status": "error", "error": "classic_failed", "elapsed_ms": 0}
|
|
vision = {"status": "error", "error": "vision_disabled", "elapsed_ms": 0}
|
|
with patch.object(server, "_classic", return_value=classic) as run_classic, patch.object(
|
|
server, "_vision", return_value=vision) as run_vision:
|
|
with TestClient(server.app) as client:
|
|
response = client.post("/api/compare", content=photo(), headers={
|
|
"content-type": "image/jpeg", "x-ocr-detail": "fine"})
|
|
invalid = client.post("/api/compare", content=photo(), headers={
|
|
"content-type": "image/jpeg", "x-ocr-detail": "unknown"})
|
|
self.assertEqual(response.status_code, 200)
|
|
self.assertTrue(run_classic.call_args.kwargs["fine_text"])
|
|
self.assertEqual(run_vision.call_count, 1)
|
|
self.assertEqual(invalid.status_code, 422)
|
|
self.assertEqual(invalid.json()["error"], "invalid_detail")
|
|
|
|
def test_classic_worker_uses_requested_detector_resolution(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
for name in ("PP-OCRv5_mobile_det", "latin_PP-OCRv5_mobile_rec"):
|
|
model = root / name
|
|
model.mkdir()
|
|
(model / "inference.pdiparams").touch()
|
|
seen = []
|
|
|
|
class FakePaddleOCR:
|
|
def __init__(self, **kwargs):
|
|
seen.append(kwargs)
|
|
|
|
def predict(self, _image):
|
|
return [SimpleNamespace(json={"res": {"rec_texts": ["B"],
|
|
"rec_polys": [[[1, 1], [2, 1], [2, 2], [1, 2]]]}})]
|
|
|
|
for options, expected in (([], 960), (["--fine-text"], 2048)):
|
|
output = root / "result.json"
|
|
with patch.object(classic_worker, "MODELS", root), patch.object(
|
|
sys, "argv", ["classic_worker", str(root / "input.jpg"), str(output), *options]), patch.dict(
|
|
sys.modules, {"paddleocr": SimpleNamespace(PaddleOCR=FakePaddleOCR)}):
|
|
classic_worker.main()
|
|
self.assertEqual(seen[-1]["text_det_limit_side_len"], expected)
|
|
self.assertEqual(seen[-1]["text_det_limit_type"], "max")
|
|
self.assertEqual(json.loads(output.read_text())["lines"][0]["text"], "B")
|
|
|
|
def test_worker_launch_errors_are_reported_per_engine(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
model_home = Path(directory)
|
|
for name in ("PP-OCRv5_mobile_det", "latin_PP-OCRv5_mobile_rec"):
|
|
model = model_home / "official_models" / name
|
|
model.mkdir(parents=True)
|
|
(model / "inference.pdiparams").touch()
|
|
with patch.object(server, "MODEL_HOME", model_home), patch.object(
|
|
server.subprocess, "run", side_effect=OSError("worker unavailable")):
|
|
classic = server._classic(photo(), model_home)
|
|
self.assertEqual(classic["error"], "classic_failed")
|
|
with patch.dict(server.os.environ, {"OCR_VISION_MODE": "ssh",
|
|
"OCR_VISION_SSH_TARGET": "test@example.invalid",
|
|
"OCR_VISION_REMOTE_DIR": "/tmp/test"}):
|
|
with patch.object(server.subprocess, "run", side_effect=OSError("ssh unavailable")):
|
|
vision = server._vision(photo())
|
|
self.assertEqual(vision["error"], "vision_failed")
|
|
|
|
def test_invalid_input_is_rejected_before_inference(self):
|
|
with TestClient(server.app) as client:
|
|
response = client.post("/api/compare", content=b"not image", headers={"content-type": "image/jpeg"})
|
|
self.assertEqual(response.status_code, 422)
|
|
self.assertEqual(response.json()["error"], "invalid_image")
|
|
|
|
def test_vision_table_keeps_all_rows_with_region_evidence(self):
|
|
payload = {"width": 900, "height": 600, "parsing_res_list": [{
|
|
"block_label": "table", "block_bbox": [20, 20, 850, 500],
|
|
"block_content": "<table><tr><td>D.3</td><td>SAMPLE 123</td></tr>" +
|
|
"<tr><td>P.3</td><td>BENZIN</td></tr></table>",
|
|
}]}
|
|
result = extract(payload)
|
|
self.assertEqual(len(result["lines"]), 2)
|
|
fields = {f["code"]: f for f in assign_fields(result["lines"], allow_adjacent=False)}
|
|
self.assertEqual(fields["P.3"]["value"], "BENZIN")
|
|
self.assertEqual(fields["D.3"]["box"], [20, 20, 850, 500])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|