label_processing.py

"""
Label detection and image preparation.

Turns a PDF page into a print-ready 1-bit label image. Finding the label uses
two strategies, best first:

1. PDF structure: most carrier labels (UPS, FedEx, Amazon, eBay, ...) are a
   single embedded image, so its exact position can be read from the PDF.
2. Raster fallback: render the page, group nearby ink into blocks, and keep
   the largest block. A dashed "cut here" border around the label is peeled
   off so only the label itself is printed.
"""
import logging
from collections import deque
from typing import Optional

import fitz  # PyMuPDF
from PIL import Image, ImageFilter

logger = logging.getLogger(__name__)

RENDER_DPI = 300      # Render resolution before scaling to the printer
MANUAL_CROP_DPI = 150  # Resolution of the full-page view used for manual crop selection

# An embedded image counts as "the label" when it is at least this big...
MIN_LABEL_SIDE_PT = 144          # 2 inches
MIN_LABEL_PAGE_FRACTION = 0.10
# ...but not when it covers the whole page (a scanned page, not a label).
MAX_LABEL_PAGE_FRACTION = 0.90
# Pages up to this size are treated as being the label itself (4x6" plus slack).
MAX_LABEL_PAGE_SHORT_PT = 4.5 * 72
MAX_LABEL_PAGE_LONG_PT = 6.5 * 72


def is_label_sized(page: "fitz.Page") -> bool:
    """True when the page is no bigger than a 4x6 label (with some slack), not a full sheet."""
    short, long = sorted((page.rect.width, page.rect.height))
    return short <= MAX_LABEL_PAGE_SHORT_PT and long <= MAX_LABEL_PAGE_LONG_PT


def find_label_rect(page: "fitz.Page") -> Optional["fitz.Rect"]:
    """
    Find the label on a page from the PDF's own structure.

    Returns the bounding box (in PDF points) of the largest label-sized
    embedded image, or None if the page has no such image.
    """
    page_area = page.rect.width * page.rect.height
    best = None
    best_area = 0.0

    for info in page.get_image_info():
        rect = fitz.Rect(info["bbox"]) & page.rect
        if rect.is_empty:
            continue
        area = rect.width * rect.height
        fraction = area / page_area
        if min(rect.width, rect.height) < MIN_LABEL_SIDE_PT:
            continue
        if not MIN_LABEL_PAGE_FRACTION <= fraction <= MAX_LABEL_PAGE_FRACTION:
            continue
        if area > best_area:
            best, best_area = rect, area

    if best is not None:
        logger.info(f"Label image found on page at {tuple(round(v, 1) for v in best)}")
    return best


def pick_label_page(doc: "fitz.Document") -> int:
    """Return the index of the page most likely to hold the label (first page by default)."""
    best_index, best_area = 0, 0.0
    for index, page in enumerate(doc):
        rect = find_label_rect(page)
        if rect is not None and rect.width * rect.height > best_area:
            best_index, best_area = index, rect.width * rect.height
    return best_index


def render_page(page: "fitz.Page", dpi: int = RENDER_DPI, clip: Optional["fitz.Rect"] = None) -> Image.Image:
    """Render a page (or part of one) to a grayscale PIL image."""
    pix = page.get_pixmap(dpi=dpi, clip=clip, colorspace=fitz.csGRAY, alpha=False)
    return Image.frombytes("L", (pix.width, pix.height), pix.samples)


def trim_whitespace(image: Image.Image, threshold: int = 245, margin: int = 10) -> Image.Image:
    """Crop away white space around the content, keeping a small margin."""
    mask = image.point(lambda v: 255 if v < threshold else 0)
    bbox = mask.getbbox()
    if not bbox:
        return image
    left, top, right, bottom = bbox
    return image.crop((
        max(0, left - margin),
        max(0, top - margin),
        min(image.width, right + margin),
        min(image.height, bottom + margin),
    ))


def _ink_blocks(mask: Image.Image) -> list:
    """
    Group dark pixels of a small binary mask into connected blocks.

    Returns a list of bounding boxes (left, top, right, bottom), right/bottom exclusive.
    """
    width, height = mask.size
    data = bytearray(mask.tobytes())  # 255 = ink, 0 = background; visited pixels are zeroed
    blocks = []

    for start, value in enumerate(data):
        if not value:
            continue
        data[start] = 0
        queue = deque([start])
        left = right = start % width
        top = bottom = start // width
        while queue:
            index = queue.popleft()
            x, y = index % width, index // width
            left, right = min(left, x), max(right, x)
            top, bottom = min(top, y), max(bottom, y)
            for neighbor, ok in (
                (index - 1, x > 0),
                (index + 1, x < width - 1),
                (index - width, y > 0),
                (index + width, y < height - 1),
            ):
                if ok and data[neighbor]:
                    data[neighbor] = 0
                    queue.append(neighbor)
        blocks.append((left, top, right + 1, bottom + 1))

    return blocks


def _row_profile(image: Image.Image) -> list:
    """Average brightness of each row (255 = completely white)."""
    return list(image.resize((1, image.height), Image.Resampling.BOX).getdata())


def _column_profile(image: Image.Image) -> list:
    """Average brightness of each column (255 = completely white)."""
    return list(image.resize((image.width, 1), Image.Resampling.BOX).getdata())


def _frame_edge(profile: list, min_gap: int, max_line: int) -> Optional[int]:
    """
    Check whether a profile starts with a thin border line followed by a white gap.

    Returns the index where the content after the gap begins, or None if the
    edge does not look like a separate border.
    """
    is_ink = [v < 250 for v in profile]
    i = 0
    while i < len(is_ink) and not is_ink[i]:
        i += 1
    line_start = i
    while i < len(is_ink) and is_ink[i]:
        i += 1
    if i == line_start or i - line_start > max_line:
        return None
    gap_start = i
    while i < len(is_ink) and not is_ink[i]:
        i += 1
    if i - gap_start < min_gap or i >= len(is_ink):
        return None
    return i


def _peel_frame(image: Image.Image) -> Image.Image:
    """
    Remove a "cut along this line" border that sits apart from the label.

    Only peels when all four sides show a thin line followed by white space,
    so a label's own tight border is left alone.
    """
    width, height = image.size
    max_line = max(4, int(min(width, height) * 0.02))
    min_gap = max(3, int(min(width, height) * 0.005))

    rows = _row_profile(image)
    top = _frame_edge(rows, min_gap, max_line)
    bottom = _frame_edge(rows[::-1], min_gap, max_line)
    if top is None or bottom is None:
        return image
    inner = image.crop((0, top, width, height - bottom))

    cols = _column_profile(inner)
    left = _frame_edge(cols, min_gap, max_line)
    right = _frame_edge(cols[::-1], min_gap, max_line)
    if left is None or right is None:
        return image

    logger.info("Removed cut-line border around label")
    return inner.crop((left, 0, inner.width - right, inner.height))


def crop_to_label(image: Image.Image) -> Image.Image:
    """
    Find the label in a rendered page image and crop to it.

    Ink is grouped into blocks (paragraphs, the label, a border box, ...) on a
    downscaled copy; the block with the largest footprint is taken as the label.
    """
    if image.mode != "L":
        image = image.convert("L")

    # Work at roughly 40 DPI: fast, and close-together label parts merge into one block.
    scale = 400 / max(image.size)
    small = image.resize(
        (max(1, int(image.width * scale)), max(1, int(image.height * scale))),
        Image.Resampling.BOX,
    )
    mask = small.point(lambda v: 255 if v < 220 else 0).filter(ImageFilter.MaxFilter(5))

    blocks = _ink_blocks(mask)
    if not blocks:
        return image

    left, top, right, bottom = max(blocks, key=lambda b: (b[2] - b[0]) * (b[3] - b[1]))
    pad = 2  # small-image pixels, to cover rounding from the downscale
    box = (
        max(0, int((left - pad) / scale)),
        max(0, int((top - pad) / scale)),
        min(image.width, int((right + pad) / scale)),
        min(image.height, int((bottom + pad) / scale)),
    )
    logger.info(f"Largest ink block: {box} of {image.size}")

    label = trim_whitespace(image.crop(box), margin=0)
    label = _peel_frame(label)
    return trim_whitespace(label)


def render_label(
    page: "fitz.Page",
    auto_crop: bool = True,
    manual_crop: Optional[tuple] = None,
    dpi: int = RENDER_DPI,
) -> Image.Image:
    """
    Render the part of the page to print, as a grayscale image.

    manual_crop is (left, top, right, bottom) in pixels of the page rendered at
    MANUAL_CROP_DPI, as selected in the app's preview.
    """
    if manual_crop:
        factor = 72.0 / MANUAL_CROP_DPI
        clip = fitz.Rect(*(v * factor for v in manual_crop)) & page.rect
        if not clip.is_empty:
            return render_page(page, dpi, clip)

    if not auto_crop:
        return render_page(page, dpi)

    if is_label_sized(page):
        # The page is already the label (e.g. a 4x6 PDF); just trim the margins
        return trim_whitespace(render_page(page, dpi))

    rect = find_label_rect(page)
    if rect is not None:
        return trim_whitespace(render_page(page, dpi, rect))

    logger.info("No embedded label image, detecting label from page content")
    return crop_to_label(render_page(page, dpi))


def fit_to_label(
    image: Image.Image,
    label_width_mm: float,
    label_height_mm: float,
    printer_dpi: int = 203,
    auto_rotate: bool = True,
    flip_vertical: bool = False,
) -> Image.Image:
    """
    Rotate and scale an image to fill the label, returning a 1-bit image
    (black = print) of exactly the label's size in printer dots.
    """
    target_w = int(label_width_mm * printer_dpi / 25.4)
    target_h = int(label_height_mm * printer_dpi / 25.4)

    if auto_rotate and (image.width > image.height) != (target_w > target_h):
        # Sideways labels almost always read bottom-to-top, so turn clockwise
        image = image.rotate(-90, expand=True)

    if flip_vertical:
        # 180 degrees turns it upside down without mirroring the text
        image = image.rotate(180)

    scale = min(target_w / image.width, target_h / image.height)
    new_size = (max(1, int(image.width * scale)), max(1, int(image.height * scale)))
    image = image.resize(new_size, Image.Resampling.LANCZOS)

    final = Image.new("L", (target_w, target_h), 255)
    final.paste(image, ((target_w - new_size[0]) // 2, (target_h - new_size[1]) // 2))
    return final.point(lambda v: 0 if v < 128 else 255, "1")


def prepare_label(
    pdf_path,
    page_index: int = 0,
    label_width_mm: float = 101,
    label_height_mm: float = 152,
    auto_crop: bool = True,
    auto_rotate: bool = True,
    flip_vertical: bool = False,
    manual_crop: Optional[tuple] = None,
    printer_dpi: int = 203,
) -> Image.Image:
    """Full pipeline: PDF page -> print-ready 1-bit label image (not inverted)."""
    with fitz.open(pdf_path) as doc:
        if len(doc) == 0:
            raise ValueError("PDF has no pages")
        page = doc[max(0, min(page_index, len(doc) - 1))]
        image = render_label(page, auto_crop=auto_crop, manual_crop=manual_crop)
    return fit_to_label(image, label_width_mm, label_height_mm, printer_dpi, auto_rotate, flip_vertical)