Spaces:
Running on Zero
Running on Zero
File size: 3,126 Bytes
9d5e109 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | 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)
|