tests/test_image_transforms.py

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