preprocess.py

"""Line-level image preprocessing before TrOCR inference."""
import cv2, numpy as np
from typing import List, Tuple

def decode_image(data: bytes):
    arr = np.frombuffer(data, np.uint8)
    return cv2.imdecode(arr, cv2.IMREAD_COLOR)

def encode_png(img: np.ndarray) -> bytes:
    _, buf = cv2.imencode(".png", img)
    return buf.tobytes()

def segment_lines(image: np.ndarray, min_height: int = 8) -> List[Tuple[int, np.ndarray]]:
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if len(image.shape) == 3 else image
    _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    proj = np.sum(binary, axis=1)
    in_line, start, lines = False, 0, []
    for i, v in enumerate(proj):
        if v > 0 and not in_line:
            in_line, start = True, i
        elif v == 0 and in_line:
            in_line = False
            crop = gray[max(0, start-2):i+2, :]
            if crop.shape[0] >= min_height:
                lines.append((len(lines), crop))
    if in_line:
        crop = gray[max(0, start-2):, :]
        if crop.shape[0] >= min_height:
            lines.append((len(lines), crop))
    return lines if lines else [(0, gray)]