-
Notifications
You must be signed in to change notification settings - Fork 58
/
detect.py
42 lines (37 loc) · 1.62 KB
/
detect.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
import argparse
import matplotlib.pyplot as plt
from lib.text_detection import TextDetection
from lib.utils import plt_show
from lib.config import Config
## Arguments
parser = argparse.ArgumentParser()
parser.add_argument('-i', '--input', type=str, required=True, help='Path to the input image')
parser.add_argument('-o', '--output', type=str, help='Path to the output image')
parser.add_argument('-s', '--show', type=str, help='Show outputs')
parser.add_argument('-d', '--direction', default='both+', type=str,
choices=set(('light', 'dark', 'both', 'both+')), help='Text searching')
parser.add_argument('-t', '--tesseract', action='store_true', help='Tesseract assistance')
parser.add_argument('--details', action='store_true', help='Detailed run with intermediate steps')
parser.add_argument('-f', '--fulltesseract', action='store_true', help='Full Tesseract')
args = vars(parser.parse_args())
IMAGE_FILE = args['input']
OUTPUT_FILE = args['output']
SHOW = args['show']
DIRECTION = args['direction']
TESS = args['tesseract']
DETAILS = args['details']
FULL_OCR = args['fulltesseract']
if __name__ == '__main__':
config = Config()
td = TextDetection(IMAGE_FILE, config, direction=DIRECTION, use_tesseract=TESS, details=DETAILS)
if FULL_OCR:
bounded, res = td.full_OCR()
if SHOW:
plt_show((td.img, 'Original'), (bounded, 'Final'), (res, 'Mask'))
else:
res = td.detect()
if SHOW:
plt_show((td.img, 'Original'), (td.final, 'Final'), (res, 'Mask'))
if OUTPUT_FILE:
plt.imsave(OUTPUT_FILE, td.final)
print(f'{OUTPUT_FILE} written')