Spaces:
Sleeping
Sleeping
File size: 4,625 Bytes
649ee03 | 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 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | from PIL import Image
from document_pii_redactor.entities import PIIEntity
from document_pii_redactor.image.transforms import (
apply_anonymize, apply_deidentify, group_text_entities,
)
from document_pii_redactor.pseudonym import PseudonymMapping
def _word(cat, bbox, txt, kind="text"):
return PIIEntity(category=cat, kind=kind, bbox=bbox,
l1="person", text=txt, score=0.9)
# --------------------------------------------------------------- grouping --- #
def test_adjacent_same_category_words_merge():
entities = [
_word("primary_subject_name", (10, 10, 50, 30), "John"),
_word("primary_subject_name", (55, 10, 90, 30), "Doe"),
]
groups = group_text_entities(entities)
assert len(groups) == 1
assert groups[0].text == "John Doe"
assert groups[0].bbox == (10, 10, 90, 30)
def test_different_lines_do_not_merge():
entities = [
_word("primary_subject_name", (10, 10, 50, 30), "John"),
_word("primary_subject_name", (10, 50, 50, 70), "Asha"),
]
assert len(group_text_entities(entities)) == 2
def test_large_horizontal_gap_does_not_merge():
entities = [
_word("primary_subject_name", (10, 10, 50, 30), "John"),
_word("primary_subject_name", (200, 10, 240, 30), "Asha"),
]
assert len(group_text_entities(entities)) == 2
def test_different_categories_do_not_merge():
entities = [
_word("primary_subject_name", (10, 10, 50, 30), "John"),
_word("city_district", (55, 10, 90, 30), "Pune"),
]
assert len(group_text_entities(entities)) == 2
def test_visual_entities_are_not_grouped():
entities = [_word("signature", (10, 10, 80, 40), None, kind="visual")]
assert group_text_entities(entities) == []
def test_text_inside_visual_entity_is_dropped():
# A logo's own lettering gets OCR'd as a text word inside the logo box;
# the placeholder already covers it — no label should render there.
entities = [
_word("logo", (100, 100, 300, 250), None, kind="visual"),
_word("brandname", (150, 140, 260, 200), "WS"),
_word("brandname", (400, 140, 520, 200), "Clinic"), # outside: kept
]
groups = group_text_entities(entities)
assert [g.text for g in groups] == ["Clinic"]
def test_mixed_box_heights_on_one_printed_line_still_merge():
# Real OCR gives differing box heights for words on the same line; that
# must not split "John Doe" into Person_1 + Person_2.
entities = [
_word("primary_subject_name", (10, 10, 50, 30), "John"),
_word("primary_subject_name", (55, 16, 90, 28), "Doe"),
]
groups = group_text_entities(entities)
assert len(groups) == 1
assert groups[0].text == "John Doe"
def test_groups_come_back_in_reading_order():
# Input deliberately lists the lower line first; numbering follows the
# order groups are returned, so reading order matters.
entities = [
_word("primary_subject_name", (10, 50, 50, 70), "Asha"),
_word("primary_subject_name", (10, 10, 50, 30), "John"),
]
groups = group_text_entities(entities)
assert [g.text for g in groups] == ["John", "Asha"]
def test_chain_survives_interleaved_other_category_word():
# "14 MG Road, <Bangalore,> Karnataka" — the address run continues across
# the city word instead of restarting a new numbered group.
entities = [
_word("street_address", (0, 0, 40, 20), "14"),
_word("city_district", (44, 0, 58, 20), "Blr"),
_word("street_address", (62, 0, 100, 20), "Road"),
]
groups = group_text_entities(entities)
assert len(groups) == 2
street = next(g for g in groups if g.category == "street_address")
assert street.text == "14 Road"
# --------------------------------------------------------------- appliers --- #
def test_deidentify_assigns_one_pseudonym_per_merged_entity():
img = Image.new("RGB", (200, 60), "white")
entities = [
_word("primary_subject_name", (10, 10, 50, 30), "John"),
_word("primary_subject_name", (55, 10, 90, 30), "Doe"),
]
mapping = PseudonymMapping()
out = apply_deidentify(img, entities, mapping)
assert out.size == img.size
assert mapping.entries["Person"] == {"Person_1": "John Doe"}
def test_anonymize_blacks_out_visual_entities_and_returns_copy():
img = Image.new("RGB", (100, 100), "white")
entities = [_word("signature", (20, 20, 80, 80), None, kind="visual")]
out = apply_anonymize(img, entities)
assert out.getpixel((50, 50)) == (0, 0, 0)
assert img.getpixel((50, 50)) == (255, 255, 255) # input untouched
|