Commit 5f154f0a authored by dmMaze's avatar dmMaze
Browse files

make merge lines optional for ysgyolo, close #866

parent a515a420

modules/textdetector/detector_ysg.py

+10 −2
Original line number Diff line number Diff line
@@ -7,7 +7,7 @@ import numpy as np 
import cv2



from .base import register_textdetectors, TextDetectorBase, TextBlock, DEVICE_SELECTOR

from utils.textblock import mit_merge_textlines, sort_regions

from utils.textblock import mit_merge_textlines, sort_regions, examine_textblk

from utils.textblock_mask import canny_flood

from utils.split_text_region import manga_split, split_textblock

from utils.imgproc_utils import xywh2xyxypoly
@@ -50,6 +50,7 @@ class YSGYoloDetector(TextDetectorBase): 
            'path_filter': '*.pt *.ckpt *.pth *.safetensors',

            'size': 'median'

        },

        'merge text lines': True,

        'confidence threshold': 0.3,

        'IoU threshold': 0.5,

        'font size multiplier': 1.,
@@ -199,7 +200,14 @@ class YSGYoloDetector(TextDetectorBase): 
                    cv2.fillPoly(mask, [pts], 255)

                pts_list += xyxy_list.tolist()



        if self.get_param_value('merge text lines'):

            blk_list += mit_merge_textlines(pts_list, width=im_w, height=im_h)

        else:

            for pts in pts_list:

                blk = TextBlock(lines=[pts])

                blk.adjust_bbox()

                examine_textblk(blk, im_w, im_h)

                blk_list.append(blk)

        blk_list = sort_regions(blk_list)



        fnt_rsz = self.get_param_value('font size multiplier')