| 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() |
|
|
