main.py 2.3 KB
from pytesseract import image_to_string
import argparse
import cv2


def main():
    parser = argparse.ArgumentParser(description='OCR program')
    parser.add_argument('img_file', type=str, help="Input Image File Name")
    parser.add_argument('--op', type=int,
                        help="Input option number to 0 ~12", default=0)
    parser.add_argument('--engine', type=str,
                        help="Select engine to Google-Vision and Tessreact")
    args = parser.parse_args()

    file_path = 'C:\\Users\\argos\\Desktop\\' + args.img_file
    img = cv2.imread(file_path)
    select_freset(args.op, img)


def tesseract_orc_to_file(img):
    text = image_to_string(img, lang='kor')
    print(text) # debug
    with open('result\\foo.txt', "w") as f:
        f.write(text)


def select_freset(type, img):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    if type == 1:
        #gray = cv2.GaussianBlur(gray, ksize=(3, 3), sigmaX=0)
        gray = cv2.bilateralFilter(gray, 9, 75, 75)
        _, img = cv2.threshold(gray, 30, 255, cv2.THRESH_BINARY)
    elif type == 2:
        gray = cv2.bilateralFilter(gray, 9, 75, 75)
        _, img = cv2.threshold(gray, 30, 255, cv2.THRESH_BINARY_INV)
    elif type == 3:
        gray = cv2.bilateralFilter(gray, 9, 75, 75)
        _, img = cv2.threshold(gray, 225, 255, cv2.THRESH_BINARY)
    elif type == 4:
        gray = cv2.bilateralFilter(gray, 9, 75, 75)
        _, img = cv2.threshold(gray, 30, 225, cv2.THRESH_BINARY_INV)
    elif type == 5:
        _, img = cv2.threshold(gray, 30, 255, cv2.THRESH_BINARY)
    elif type == 6:
        _, img = cv2.threshold(gray, 30, 255, cv2.THRESH_BINARY_INV)
    elif type == 7:
        _, img = cv2.threshold(gray, 225, 255, cv2.THRESH_BINARY)
    elif type == 8:
        _, img = cv2.threshold(gray, 225, 255, cv2.THRESH_BINARY_INV)
    elif type == 9:
        _, img = cv2.threshold(gray, 98, 255, cv2.THRESH_BINARY)
    elif type == 10:
        _, img = cv2.threshold(gray, 98, 255, cv2.THRESH_BINARY_INV)
    elif type == 11:
        gray = cv2.bilateralFilter(gray, 9, 75, 75)
        _, img = cv2.threshold(gray, 98, 255, cv2.THRESH_BINARY)
    elif type == 12:
        gray = cv2.bilateralFilter(gray, 9, 75, 75)
        _, img = cv2.threshold(gray, 98, 255, cv2.THRESH_BINARY_INV)
    tesseract_orc_to_file(img)


if __name__ == "__main__":
    main()