added cropped text region export

file_utils.py

@@ -1,7 +1,8 @@
 # -*- coding: utf-8 -*-
 import os
-import numpy as np
 import cv2
+import copy
+import numpy as np
 
 
 def create_dir(_dir):
@@ -12,7 +13,6 @@ def create_dir(_dir):
         os.makedirs(_dir)
 
 
-# borrowed from https://github.com/lengstrom/fast-style-transfer/blob/master/src/utils.py
 def get_files(img_dir):
     imgs, masks, xmls = list_files(img_dir)
     return imgs, masks, xmls
@@ -41,30 +41,69 @@ def list_files(in_path):
     return img_files, mask_files, gt_files
 
 
-def saveResult(img_file,
-               img,
-               boxes,
-               dirname='./result/',
-               verticals=None,
-               texts=None):
+def export_detected_region(image, points, file_path):
+    # points should have 1*4*2  shape
+    if len(points.shape) == 2:
+        points = np.array([np.array(points).astype(np.int32)])
+
+    # create mask with shape of image
+    mask = np.zeros(image.shape[0:2], dtype=np.uint8)
+
+    # method 1 smooth region
+    cv2.drawContours(mask, [points], -1, (255, 255, 255), -1, cv2.LINE_AA)
+
+    # method 2 not so smooth region
+    # cv2.fillPoly(mask, points, (255))
+
+    res = cv2.bitwise_and(image, image, mask=mask)
+    rect = cv2.boundingRect(points)  # returns (x,y,w,h) of the rect
+    cropped = res[rect[1]: rect[1] + rect[3], rect[0]: rect[0] + rect[2]]
+
+    # export corpped region
+    cv2.imwrite(file_path, cropped)
+
+
+def export_detected_regions(image_path, image, polys,
+                            output_dir: str = "output/"):
+    # deepcopy image so that original is not altered
+    image = copy.deepcopy(image)
+
+    # get file name
+    file_name, file_ext = os.path.splitext(os.path.basename(image_path))
+
+    # create crops dir
+    crops_dir = os.path.join(output_dir, file_name + "_crops")
+    create_dir(crops_dir)
+
+    for ind, poly in enumerate(polys):
+        file_path = os.path.join(crops_dir, "crop_" + str(ind) + ".png")
+        export_detected_region(image, points=poly, file_path=file_path)
+
+
+def export_extra_results(image_path,
+                         image,
+                         boxes,
+                         output_dir='output/',
+                         verticals=None,
+                         texts=None):
     """ save text detection result one by one
     Args:
-        img_file (str): image file name
-        img (array): raw image context
+        image_path (str): image file name
+        image (array): raw image context
         boxes (array): array of result file
             Shape: [num_detections, 4] for BB output / [num_detections, 4]
             for QUAD output
     Return:
         None
     """
-    img = np.array(img)
+    image = np.array(image)
 
     # make result file list
-    filename, file_ext = os.path.splitext(os.path.basename(img_file))
+    filename, file_ext = os.path.splitext(os.path.basename(image_path))
 
     # result directory
-    res_file = dirname + "res_" + filename + '.txt'
-    res_img_file = dirname + "res_" + filename + '.jpg'
+    res_file = output_dir + "res_" + filename + '.txt'
+    res_img_file = output_dir + "res_" + filename + '.jpg'
 
     with open(res_file, 'w') as f:
         for i, box in enumerate(boxes):
@@ -73,7 +112,7 @@ def saveResult(img_file,
             f.write(strResult)
 
             poly = poly.reshape(-1, 2)
-            cv2.polylines(img,
+            cv2.polylines(image,
                           [poly.reshape((-1, 1, 2))],
                           True,
                           color=(0, 0, 255),
@@ -86,13 +125,13 @@ def saveResult(img_file,
             if texts is not None:
                 font = cv2.FONT_HERSHEY_SIMPLEX
                 font_scale = 0.5
-                cv2.putText(img, "{}".format(texts[i]),
+                cv2.putText(image, "{}".format(texts[i]),
                             (poly[0][0]+1, poly[0][1]+1),
                             font,
                             font_scale,
                             (0, 0, 0),
                             thickness=1)
-                cv2.putText(img,
+                cv2.putText(image,
                             "{}".format(texts[i]),
                             tuple(poly[0]),
                             font,
@@ -101,4 +140,4 @@ def saveResult(img_file,
                             thickness=1)
 
     # Save result image
-    cv2.imwrite(res_img_file, img)
+    cv2.imwrite(res_img_file, image)

predict.py

@@ -1,9 +1,3 @@
-"""
-Copyright (c) 2019-present NAVER Corp.
-MIT License
-"""
-
-# -*- coding: utf-8 -*-
 import os
 import time
 
@@ -154,7 +148,8 @@ def detect_text(image_path: str,
                 mag_ratio: float = 1.5,
                 poly: bool = False,
                 show_time: bool = False,
-                refiner: bool = False):
+                refiner: bool = False,
+                export_extra: bool = True):
     """
     Arguments:
         image_path: path to the image to be processed
@@ -168,6 +163,7 @@ def detect_text(image_path: str,
         poly: enable polygon type
         show_time: show processing time
         refiner: enable link refiner
+        export_extra: export score map, detection points, box visualization
     """
 
     # create output dir
@@ -194,14 +190,19 @@ def detect_text(image_path: str,
                                                poly,
                                                show_time)
 
-    # save score text
-    filename, file_ext = os.path.splitext(os.path.basename(image_path))
-    mask_file = os.path.join(output_dir, "res_" + filename + '_mask.jpg')
-    cv2.imwrite(mask_file, score_text)
+    # export detected text regions
+    file_utils.export_detected_regions(image_path, image, polys, output_dir)
+
+    if export_extra:
+        # export score map
+        filename, file_ext = os.path.splitext(os.path.basename(image_path))
+        mask_file = os.path.join(output_dir, "res_" + filename + '_mask.jpg')
+        cv2.imwrite(mask_file, score_text)
 
-    file_utils.saveResult(image_path,
-                          image[:, :, ::-1],
-                          polys,
-                          dirname=output_dir)
+        # export detected points and box visualization
+        file_utils.export_extra_results(image_path,
+                                        image[:, :, ::-1],
+                                        polys,
+                                        output_dir=output_dir)
 
     print("elapsed time : {}s".format(time.time() - t))