deeptextdetection.py 1.29 KB
Newer Older
sghoshcvc's avatar
sghoshcvc committed
1 2 3 4
# -*- coding: utf-8 -*-
#!/usr/bin/python
import sys
import os
5
import cv2 as cv
sghoshcvc's avatar
sghoshcvc committed
6 7
import numpy as np

8 9 10 11
def main():
    print('\nDeeptextdetection.py')
    print('       A demo script of text box alogorithm of the paper:')
    print('       * Minghui Liao et al.: TextBoxes: A Fast Text Detector with a Single Deep Neural Network https://arxiv.org/abs/1611.06779\n')
sghoshcvc's avatar
sghoshcvc committed
12

13 14 15
    if (len(sys.argv) < 2):
        print(' (ERROR) You must call this script with an argument (path_to_image_to_be_processed)\n')
        quit()
sghoshcvc's avatar
sghoshcvc committed
16

17
    if not os.path.isfile('TextBoxes_icdar13.caffemodel') or not os.path.isfile('textbox.prototxt'):
18
        print " Model files not found in current directory. Aborting"
19
        print " See the documentation of text::TextDetectorCNN class to get download links."
20
        quit()
sghoshcvc's avatar
sghoshcvc committed
21

22 23
    img = cv.imread(str(sys.argv[1]))
    textSpotter = cv.text.TextDetectorCNN_create("textbox.prototxt", "TextBoxes_icdar13.caffemodel")
24
    rects, outProbs = textSpotter.detect(img);
25 26
    vis = img.copy()
    thres = 0.6
sghoshcvc's avatar
sghoshcvc committed
27

28 29 30
    for r in range(np.shape(rects)[0]):
        if outProbs[r] > thres:
            rect = rects[r]
31
            cv.rectangle(vis, (rect[0],rect[1]), (rect[0] + rect[2], rect[1] + rect[3]), (255, 0, 0), 2)
sghoshcvc's avatar
sghoshcvc committed
32

33 34
    cv.imshow("Text detection result", vis)
    cv.waitKey()
sghoshcvc's avatar
sghoshcvc committed
35

36 37
if __name__ == "__main__":
    main()