Preserve only text mask regions under 50 pixels tall

text_detection.py

@@ -6,13 +6,15 @@ TEXT_MODEL = "PP-OCRv5_server_det"
 TEXT_DEVICE = "cpu"
 TEXT_THREADS = 1
 TEXT_PADDING = 2
+TEXT_MAX_HEIGHT = 50
 
 
 def main():
     os.environ["OMP_NUM_THREADS"] = str(TEXT_THREADS)
+    import cv2
     import numpy as np
     from paddleocr import TextDetection
-    from segmentation import character_mask
+    from segmentation import MASK_VALUE, character_mask
 
     model = TextDetection(
         model_name=TEXT_MODEL, device=TEXT_DEVICE, cpu_threads=TEXT_THREADS,
@@ -29,6 +31,10 @@ def main():
     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")