Spaces:
Running on Zero
Running on Zero
| import numpy as np | |
| import pytest | |
| from PIL import Image | |
| from pixtext_ocr import WordOCR | |
| def quad(x0, y0, x1, y1): | |
| return np.float32([[x0, y0], [x1, y0], [x1, y1], [x0, y1]]) | |
| def routed_detector(parent, glyph_texts): | |
| engine = WordOCR.__new__(WordOCR) | |
| word = quad(0, 5, 150, 45) | |
| chars = [quad(x, 10, x + width, 40) for x, width in [(5, 19), (35, 24), (73, 18), (110, 30)]] | |
| engine.detect_quads = lambda *args: ([word], chars) | |
| calls = [] | |
| def recognize(crops): | |
| if not crops: | |
| return [] | |
| calls.append(len(crops)) | |
| return [(parent, 0.95)] if len(calls) == 1 else [(s, 0.9) for s in glyph_texts] | |
| engine.recognize = recognize | |
| return engine, chars, calls | |
| def test_chinese_uses_detected_glyph_boxes_and_independent_recognition(): | |
| engine, chars, calls = routed_detector('新年快乐', ['新', '年', '怪', '乐']) | |
| rows = engine.detect(Image.new('RGB', (160, 60)), 0.5, 0.2) | |
| assert calls == [1, 4] | |
| assert [r['text'] for r in rows] == ['新', '年', '怪', '乐'] | |
| assert [r['bbox'] for r in rows] == [[*q.min(axis=0), *q.max(axis=0)] for q in chars] | |
| assert all(r['level'] == 'character' and not r['enabled'] for r in rows) | |
| def test_english_remains_a_whole_word(): | |
| engine, _, calls = routed_detector('crossbody', []) | |
| rows = engine.detect(Image.new('RGB', (160, 60)), 0.5, 0.2) | |
| assert calls == [1] | |
| assert len(rows) == 1 and rows[0]['text'] == 'crossbody' and rows[0]['level'] == 'word' | |
| def test_unreadable_or_ambiguous_glyphs_remain_editable_without_phrase_labels(): | |
| engine, _, _ = routed_detector('新年快乐', ['', '新年', '快', '乐']) | |
| rows = engine.detect(Image.new('RGB', (160, 60)), 0.5, 0.2) | |
| assert len(rows) == 4 | |
| assert [r['text'] for r in rows] == ['', '', '快', '乐'] | |
| assert rows[0]['confidence'] == rows[1]['confidence'] == 0 | |
| def test_mixed_text_groups_latin_letters_but_keeps_han_individual(): | |
| engine, _, _ = routed_detector('AB中文', ['A', 'B', '中', '文']) | |
| rows = engine.detect(Image.new('RGB', (160, 60)), 0.5, 0.2) | |
| assert [(r['text'], r['level']) for r in rows] == [('AB', 'word'), ('中', 'character'), ('文', 'character')] | |
| def test_pure_chinese_does_not_merge_glyphs_misrecognized_as_latin(): | |
| engine, _, _ = routed_detector('新年快乐', ['新', 'L', 'S', '乐']) | |
| rows = engine.detect(Image.new('RGB', (160, 60)), 0.5, 0.2) | |
| assert len(rows) == 4 and all(r['level'] == 'character' for r in rows) | |
| def test_craft_affinity_joins_words_but_does_not_join_character_boxes(): | |
| pytest.importorskip('easyocr') | |
| import torch | |
| scores = torch.zeros((1, 64, 128, 2)) | |
| for x in [24, 55, 90]: | |
| scores[0, 24:40, x - 6:x + 6, 0] = 0.95 | |
| scores[0, 29:35, 24:91, 1] = 0.9 | |
| engine = WordOCR.__new__(WordOCR) | |
| engine.detector = lambda tensor: (scores, None) | |
| words, chars = engine.detect_quads(np.zeros((80, 160, 3), dtype=np.uint8), 0.5, 0.2) | |
| assert len(words) == 1 and len(chars) == 3 | |
| centers = sorted(float(q.mean(axis=0)[0]) for q in chars) | |
| assert np.allclose(centers, np.array([24, 55, 90]) / 1.5 * 2, atol=3) | |