from easyocr

 * fix typo
 * add estimate_num_chars option
 * return mapper

Author: wkpark

craft.py

@@ -1,4 +1,4 @@
-"""  
+"""
 Copyright (c) 2019-present NAVER Corp.
 MIT License
 """
@@ -54,7 +54,7 @@ def __init__(self, pretrained=False, freeze=False):
         init_weights(self.upconv3.modules())
         init_weights(self.upconv4.modules())
         init_weights(self.conv_cls.modules())
-        
+
     def forward(self, x):
         """ Base network """
         sources = self.basenet(x)
@@ -78,8 +78,3 @@ def forward(self, x):
         y = self.conv_cls(feature)
 
         return y.permute(0,2,3,1), feature
-
-if __name__ == '__main__':
-    model = CRAFT(pretrained=True).cuda()
-    output, _ = model(torch.randn(1, 3, 768, 768).cuda())
-    print(output.shape)

craft_utils.py

@@ -7,16 +7,17 @@
 import numpy as np
 import cv2
 import math
+from scipy.ndimage import label
 
-""" auxilary functions """
+""" auxiliary functions """
 # unwarp corodinates
 def warpCoord(Minv, pt):
     out = np.matmul(Minv, (pt[0], pt[1], 1))
     return np.array([out[0]/out[2], out[1]/out[2]])
-""" end of auxilary functions """
+""" end of auxiliary functions """
 
 
-def getDetBoxes_core(textmap, linkmap, text_threshold, link_threshold, low_text):
+def getDetBoxes_core(textmap, linkmap, text_threshold, link_threshold, low_text, estimate_num_chars=False):
     # prepare data
     linkmap = linkmap.copy()
     textmap = textmap.copy()
@@ -42,6 +43,12 @@ def getDetBoxes_core(textmap, linkmap, text_threshold, link_threshold, low_text)
         # make segmentation map
         segmap = np.zeros(textmap.shape, dtype=np.uint8)
         segmap[labels==k] = 255
+        if estimate_num_chars:
+            _, character_locs = cv2.threshold((textmap - linkmap) * segmap /255., text_threshold, 1, 0)
+            _, n_chars = label(character_locs)
+            mapper.append(n_chars)
+        else:
+            mapper.append(k)
         segmap[np.logical_and(link_score==1, text_score==0)] = 0   # remove link area
         x, y = stats[k, cv2.CC_STAT_LEFT], stats[k, cv2.CC_STAT_TOP]
         w, h = stats[k, cv2.CC_STAT_WIDTH], stats[k, cv2.CC_STAT_HEIGHT]
@@ -74,7 +81,6 @@ def getDetBoxes_core(textmap, linkmap, text_threshold, link_threshold, low_text)
         box = np.array(box)
 
         det.append(box)
-        mapper.append(k)
 
     return det, labels, mapper
 
@@ -160,7 +166,7 @@ def getPoly_core(boxes, labels, mapper, linkmap):
         if num_sec != 0:
             cp_section[-1] = [cp_section[-1][0] / num_sec, cp_section[-1][1] / num_sec]
 
-        # pass if num of pivots is not sufficient or segment widh is smaller than character height 
+        # pass if num of pivots is not sufficient or segment width is smaller than character height 
         if None in pp or seg_w < np.max(seg_height) * 0.25:
             polys.append(None); continue
 
@@ -224,15 +230,17 @@ def getPoly_core(boxes, labels, mapper, linkmap):
 
     return polys
 
-def getDetBoxes(textmap, linkmap, text_threshold, link_threshold, low_text, poly=False):
-    boxes, labels, mapper = getDetBoxes_core(textmap, linkmap, text_threshold, link_threshold, low_text)
+def getDetBoxes(textmap, linkmap, text_threshold, link_threshold, low_text, poly=False, estimate_num_chars=False):
+    if poly and estimate_num_chars:
+        raise Exception("Estimating the number of characters not currently supported with poly.")
+    boxes, labels, mapper = getDetBoxes_core(textmap, linkmap, text_threshold, link_threshold, low_text, estimate_num_chars)
 
     if poly:
         polys = getPoly_core(boxes, labels, mapper, linkmap)
     else:
         polys = [None] * len(boxes)
 
-    return boxes, polys
+    return boxes, polys, mapper
 
 def adjustResultCoordinates(polys, ratio_w, ratio_h, ratio_net = 2):
     if len(polys) > 0: