Merge pull request #77 from LDOUBLEV/fixocr

Fixocr

tools/infer/predict_rec.py

@@ -31,6 +31,7 @@ class TextRecognizer(object):
         image_shape = [int(v) for v in args.rec_image_shape.split(",")]
         self.rec_image_shape = image_shape
         self.character_type = args.rec_char_type
+        self.rec_batch_num = args.rec_batch_num
         char_ops_params = {}
         char_ops_params["character_type"] = args.rec_char_type
         char_ops_params["character_dict_path"] = args.rec_char_dict_path
@@ -59,8 +60,8 @@ class TextRecognizer(object):
 
     def __call__(self, img_list):
         img_num = len(img_list)
-        batch_num = 30
         rec_res = []
+        batch_num = self.rec_batch_num
         predict_time = 0
         for beg_img_no in range(0, img_num, batch_num):
             end_img_no = min(img_num, beg_img_no + batch_num)

tools/infer/predict_system.py

@@ -89,7 +89,7 @@ def sorted_boxes(dt_boxes):
         sorted boxes(array) with shape [4, 2]
     """
     num_boxes = dt_boxes.shape[0]
-    sorted_boxes = sorted(dt_boxes, key=lambda x: x[0][1])
+    sorted_boxes = sorted(dt_boxes, key=lambda x: (x[0][1], x[0][0]))
     _boxes = list(sorted_boxes)
 
     for i in range(num_boxes - 1):

tools/infer/utility.py

@@ -56,6 +56,7 @@ def parse_args():
     parser.add_argument("--rec_model_dir", type=str)
     parser.add_argument("--rec_image_shape", type=str, default="3, 32, 320")
     parser.add_argument("--rec_char_type", type=str, default='ch')
+    parser.add_argument("--rec_batch_num", type=int, default=30)
     parser.add_argument(
         "--rec_char_dict_path",
         type=str,
@@ -172,7 +173,8 @@ def draw_ocr(image, boxes, txts, scores, draw_txt=True, drop_score=0.5):
                 continue
             font = ImageFont.truetype(
                 "./doc/simfang.ttf", font_size, encoding="utf-8")
-            new_txt = str(count) + ':  ' + txt + '    ' + '%.3f' % (scores[count])
+            new_txt = str(count) + ':  ' + txt + '    ' + '%.3f' % (
+                scores[count])
             draw_txt.text(
                 (20, gap * (count + 1)), new_txt, txt_color, font=font)
             count += 1