adding option for grayscale conversion with tesseract

Author: scotthlee

hamlet/tools/image.py

@@ -124,7 +124,7 @@ def flip_image(img_name, img_dir):
     return
 
 
-def check_text(file_path, lim=3):
+def check_text(file_path, lim=3, convert_to_grayscale=False):
     """Checks an image for burned-in text.
 
     Parameters
@@ -140,6 +140,8 @@ def check_text(file_path, lim=3):
     True if the image has more than lim words, else False.
     """
     img = cv2.imread(file_path)
+    if convert_to_grayscale:
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
     tess = pytesseract.image_to_data(img, output_type=Output.DICT)
     words = tess['text']
     if np.any([len(w) > lim for w in words]):

text_removal.py

@@ -28,24 +28,28 @@
                         type=int,
                         default=3,
                         help='maximum allowable number of words per image')
+    parser.add_argument('--convert_to_grayscale',
+                        action='store_true')
     parser.add_argument('--no_multiprocessing',
                         action='store_true')
-    parser.set_defaults(no_multiprocessing=False)
+    parser.set_defaults(no_multiprocessing=False,
+                        convert_to_grayscale=False)
     args = parser.parse_args()
 
     # Setting globals
     IMG_DIR = args.img_dir
     TEXT_DIR = args.text_dir
     NUM_WORDS = args.num_words
     USE_MULTIPROCESSING =  not args.no_multiprocessing
+    GRAY = args.convert_to_grayscale
 
     # Importing the data
     files = os.listdir(IMG_DIR)
 
     # Checking the files
     if USE_MULTIPROCESSING:
         with Pool() as p:
-            input = [(IMG_DIR + f, NUM_WORDS) for f in files]
+            input = [(IMG_DIR + f, NUM_WORDS, GRAY) for f in files]
             res1 = p.starmap(check_text, input)
             p.close()
             p.join()
@@ -59,7 +63,7 @@
             p.close()
             p.join()
     else:
-        res1 = [check_text(IMG_DIR + f, NUM_WORDS) for f in files]
+        res1 = [check_text(IMG_DIR + f, NUM_WORDS, GRAY) for f in files]
         with_text = np.where(res1)[0]
         to_move = [files[i] for i in with_text]
         res2 = [os.rename(IMG_DIR + f, TEXT_DIR + f) for f in to_move]