import os
import sys
from PIL import Image

TEXT_MODEL = "PP-OCRv5_server_det"
TEXT_DEVICE = "cpu"
TEXT_THREADS = 1
TEXT_PADDING = 2
TEXT_MAX_HEIGHT = 40


def main():
    os.environ["OMP_NUM_THREADS"] = str(TEXT_THREADS)
    import cv2
    import numpy as np
    from paddleocr import TextDetection
    from segmentation import MASK_VALUE, character_mask

    model = TextDetection(
        model_name=TEXT_MODEL, device=TEXT_DEVICE, cpu_threads=TEXT_THREADS,
        enable_mkldnn=False,
    )
    if len(sys.argv) == 1:
        return

    input_path, output_path = sys.argv[1:]
    with Image.open(input_path) as image:
        rgb_image = np.asarray(image.convert("RGB"))
    bgr_image = np.ascontiguousarray(rgb_image[:, :, ::-1])
    polygons = []
    for result in model.predict(input=bgr_image):
        polygons.extend(result["dt_polys"])
    mask = character_mask(rgb_image, polygons, margin=TEXT_PADDING)
    count, labels, stats, _ = cv2.connectedComponentsWithStats(mask, connectivity=8)
    keep = np.zeros(count, dtype=np.uint8)
    keep[1:] = (stats[1:, cv2.CC_STAT_HEIGHT] < TEXT_MAX_HEIGHT) * MASK_VALUE
    mask = keep[labels]
    Image.fromarray(mask).save(output_path, format="PNG")


if __name__ == "__main__":
    main()