HYTOS

# coding: utf-8

"""

    This is text detector module

"""

import sys

import os

import cv2

import numpy as np

from PyQt5.QtCore import *

from PyQt5.QtGui import *

from PyQt5.QtWidgets import *

from PyQt5.QtSvg import *

from AppDocData import *

import TextInfo as ti

import tesseract_ocr_module as TOCR

MIN_TEXT_SIZE = 10

THREAD_MAX_WORKER = os.cpu_count()

class TextDetector:

    '''

        @brief  constructor

        @author humkyung

        @date   2018.07.11

    '''

    def __init__(self):

        self.textInfoList = []

        self.otherTextInfoList = []

    '''

        @brief  detect text areas

        @author humkyung

        @date   2018.06.16

    '''

    def detectTextAreas(self, img, offset):

        tInfoList = []

        try:

            tInfoList = self.getTextAreaInfo(img, offset[0], offset[1])

        except Exception as ex:

            print('error occured({}) in {}:{}'.format(ex, sys.exc_info()[-1].tb_frame.f_code.co_filename, sys.exc_info()[-1].tb_lineno))

        return tInfoList

    '''

        @brief      Get Text Area info by contour

        @author     Jeongwoo

        @date       2018.06.05

        @history    2018.06.08  Jeongwoo    Add angle

                    humkyung 2018.06.18 fixed logic to detect text area

    '''

    def getTextAreaInfo(self, imgGray, offsetX, offsetY):

        from AppDocData import AppDocData

        appDocData = AppDocData.instance()

        project = appDocData.getCurrentProject()

        configs = appDocData.getConfigs('Text Size', 'Max Text Size')

        maxTextSize = int(configs[0].value) if 1 == len(configs) else 100

        contourImg = np.ones(imgGray.shape, np.uint8) * 255

        binaryImg,mask = cv2.threshold(imgGray, 127, 255, cv2.THRESH_BINARY)

        image, contours, hierarchy = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)

        for contour in contours:

            # remove too big one

            [x, y, w, h] = cv2.boundingRect(contour)

            if (w > maxTextSize and h > maxTextSize): continue

            area = cv2.contourArea(contour, True)

            if area >= 0:

                cv2.drawContours(contourImg, [contour], -1, (0,0,0), -1)

                cv2.drawContours(contourImg, [contour], -1, (255,255,255), 1)

            else:

                cv2.drawContours(contourImg, [contour], -1, (255,255,255), -1)

        path = os.path.join(project.getTempPath(), 'OCR_{}.png'.format(appDocData.imgName))

        cv2.imwrite(path, contourImg)

        rects = []

        configs = appDocData.getConfigs('Text Recognition', 'Expand Size')

        expandSize = int(configs[0].value) if 1 == len(configs) else 10

        configs = appDocData.getConfigs('Text Recognition', 'Shrink Size')

        shrinkSize = int(configs[0].value) if 1 == len(configs) else 0

        eroded = cv2.erode(contourImg, np.ones((expandSize,expandSize), np.uint8))

        path = os.path.join(project.getTempPath(), 'ERODED_OCR_{}.png'.format(appDocData.imgName))

        cv2.imwrite(path, eroded)

        eroded = cv2.bitwise_not(eroded)

        image, contours, hierarchy = cv2.findContours(eroded, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE)

        for contour in contours:

            area = cv2.contourArea(contour, True)

            if area < 0:

                [x, y, w, h] = cv2.boundingRect(contour)

                if (w < 10 and h < 10) or (w > maxTextSize and h > maxTextSize): continue; # skip too small or big one

                img = contourImg[y:y+h, x:x+w]

                img = cv2.bitwise_not(img)

                horizontal = 0

                vertical = 0

                if w > maxTextSize:

                    horizontal = 1

                elif h > maxTextSize:

                    vertical = 1

                else:

                    if shrinkSize > 0:

                        img = cv2.erode(img, np.ones((shrinkSize,shrinkSize), np.uint8))

                    _, _contours, _ = cv2.findContours(img, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE)

                    for xx in _contours:

                        [_x, _y, _w, _h] = cv2.boundingRect(xx)

                        cv2.rectangle(img, (_x, _y), (_x+_w, _y+_h), 255, 1)

                        if (_w < _h) or (_w > maxTextSize and _h < maxTextSize): # width is greater than height

                            horizontal += 1 + (_w*_h)/(w*h)

                        else:

                            vertical += 1 + (_w*_h)/(w*h)

                if horizontal > vertical:

                    filePath = os.path.join(project.getTempPath(), "Tile", "H-{}-{}-{}-{}.png".format(x,y,w,h))

                else:

                    filePath = os.path.join(project.getTempPath(), "Tile", "V-{}-{}-{}-{}.png".format(x,y,w,h))

                cv2.imwrite(filePath, img)

                rects.append([0 if horizontal > vertical else 90, QRect(x, y, w, h)])

        configs = appDocData.getConfigs('Text Recognition', 'Merge Size')

        mergeSize = int(configs[0].value) if 1 == len(configs) else 10

        # merge rectangles

        intersected = True

        while intersected:

            intersected = False

            for rect in rects[:]:

                if 0 == rect[0]:

                    rectExpand = rect[1].adjusted(-mergeSize, 0, mergeSize, 0)

                else:

                    rectExpand = rect[1].adjusted(0, -mergeSize, 0, mergeSize)

                matches = [x for x in rects if (x[0] == rect[0]) and rectExpand.intersects(x[1])]

                if len(matches) > 1:

                    united = matches[0]

                    for _rect in matches:

                        united[1] = united[1].united(_rect[1])

                        if _rect in rects: rects.remove(_rect)

                    rects.append(united)

                    intersected = True

                    break

        list = []

        for rect in rects:

            angle = rect[0]

            list.append(ti.TextInfo('', round(offsetX) + rect[1].x(), round(offsetY) + rect[1].y(), rect[1].width(), rect[1].height(), angle))

        return list

    '''

        @brief      recognize text of given text info

        @author     humkyung

        @date       2018.07.24

        @history    change parameter updateProgressSignal to worker

    '''

    @staticmethod

    def recognizeTextFromImage(tInfo, imgOCR, offset, searchedSymbolList, worker, listWidget, maxProgressValue):

        import re

        res = []

        appDocData = AppDocData.instance()

        try:

            x = tInfo.getX() - round(offset[0])

            y = tInfo.getY() - round(offset[1])

            img = imgOCR[y:y+tInfo.getH(), x:x+tInfo.getW()]

            # set angle 0 if symbol contains the text area is instrumentation

            category = None

            contains = [symbol for symbol in searchedSymbolList if symbol.contains(tInfo)]

            if contains:

                _type = contains[0].getType()

                category = appDocData.getSymbolCategoryByType(_type)

                if 'Instrumentation' == category: tInfo.setAngle(0)

            # up to here

            resultTextInfo = TOCR.getTextInfo(img, (x, y), tInfo.getAngle())

            if resultTextInfo is not None and len(resultTextInfo) > 0:

                for result in resultTextInfo:

                    result.setX(result.getX() + round(offset[0]))

                    result.setY(result.getY() + round(offset[1]))

                    if 'Instrumentation' == category:

                        text = re.sub('[^a-zA-Z0-9]+', '', result.getText())

                        result.setText(text)

                res.extend(resultTextInfo)

                if listWidget is not None:

                    item = QListWidgetItem('{},{},{} is recognized'.format(resultTextInfo[0].getX(), resultTextInfo[0].getY(), resultTextInfo[0].getText()))

                    listWidget.addItem(item)

            else:

                pass

            if worker is not None: worker.updateProgress.emit(maxProgressValue, resultTextInfo[0].getText())

        except Exception as ex:

            message = 'error occured({}) in {}:{}'.format(ex, sys.exc_info()[-1].tb_frame.f_code.co_filename, sys.exc_info()[-1].tb_lineno)

            worker.displayLog.emit(MessageType.Error, message)

        return res

    '''

        @brief      read image drawing and then remove text

        @author     jwkim

        @date       

        @history    humkyung 2018.04.06 check if file exists

                    Jeongwoo 2018.05.09 Use Tesseract OCR after Azure OCR (Azure OCR : Getting text area)

                    Jeongwoo 2018.05.25 Add condition on if-statement

                    Jeongwoo 2018.06.05 Get text area data list by config.type

                    Jeongwoo 2018.06.08 Add angle Parameter on TOCR.getTextInfo

                    humkyung 2018.06.16 update proessbar while recognizing text

                    humkyung 2018.07.03 remove white space and replace given oldStr with newStr

                    humkyung 2018.07.07 change method name to recognizeText 

    '''

    def recognizeText(self, imgSrc, offset, tInfoList, searchedSymbolList, worker, listWidget, maxProgressValue, onlyTextArea = False):

        import concurrent.futures as futures

        try:

            appDocData = AppDocData.instance()

            project = appDocData.getCurrentProject()

            path = os.path.join(project.getTempPath(), 'OCR_{}.png'.format(appDocData.imgName))

            if os.path.isfile(path):

                imgOCR = cv2.imread(path, 1)

                imgOCR = cv2.threshold(cv2.cvtColor(imgOCR, cv2.COLOR_BGR2GRAY), 127, 255, cv2.THRESH_BINARY)[1]

                pool = futures.ThreadPoolExecutor(max_workers = THREAD_MAX_WORKER)

                for tInfo in tInfoList:

                    future = pool.submit(TextDetector.recognizeTextFromImage, tInfo, imgOCR, offset, searchedSymbolList, worker, listWidget, maxProgressValue)

                    data = future.result()

                    if data: self.textInfoList.extend(data)

                pool.shutdown(wait = True)

            if onlyTextArea:

                return

            # parse texts in area except Drawing area

            for area in appDocData.getAreaList():

                if area.name == 'Drawing': continue

                if area.name == 'Unit':

                    img = imgSrc[round(area.y):round(area.y+area.height), round(area.x):round(area.x+area.width)]

                    texts = TOCR.getTextInfo(img, (area.x, area.y), 0)

                    if texts is not None and len(texts) > 0:

                        appDocData.activeDrawing.setAttr('Unit', texts[0].getText())

                        self.otherTextInfoList.append([area.name, texts])

                else:

                    if area is not None and hasattr(area, 'img') and area.img is not None:

                        texts = TOCR.getTextInfo(area.img, (area.x, area.y))

                        self.otherTextInfoList.append([area.name, texts])

            if worker is not None: worker.updateProgress.emit(maxProgressValue, None)

        except Exception as ex:

            message = 'error occured({}) in {}:{}'.format(ex, sys.exc_info()[-1].tb_frame.f_code.co_filename, sys.exc_info()[-1].tb_lineno)

            worker.displayLog.emit(MessageType.Error, message)

    '''

        @brief      remove text from image

        @author     humkyung

        @date       2018.07.24

    '''

    def removeTextFromImage(self, imgSrc, offset):

        appDocData = AppDocData.instance()

        project = appDocData.getCurrentProject()

        path = os.path.join(project.getTempPath(), 'OCR_{}.png'.format(appDocData.imgName))

        if os.path.isfile(path):

            imgOCR = cv2.imread(path, 1)

            imgOCR = cv2.threshold(cv2.cvtColor(imgOCR, cv2.COLOR_BGR2GRAY), 127, 255, cv2.THRESH_BINARY)[1]

            # remove recognized text from image

            for text in self.textInfoList:

                x = round(text.getX() - offset[0])

                y = round(text.getY() - offset[1])

                width = round(text.getW())

                height = round(text.getH())

                self.removeText(imgSrc, (round(text.getX()), round(text.getY())), imgOCR[y:y+height, x:x+width])

            # up to here

    '''

        @brief  detect text in given area with givn angle

        @author humkyung

        @date   2018.07.11

    '''

    def recognizeTextInArea(self, imgSrc, area, angle=0):

        try:

            img = imgSrc[round(area[1]):round(area[1] + area[3]), round(area[0]):round(area[0] + area[2])]

            return TOCR.getTextInfo(img, (area[0], area[1]))

        except Exception as ex:

            print('error occured({}) in {}:{}'.format(ex, sys.exc_info()[-1].tb_frame.f_code.co_filename, sys.exc_info()[-1].tb_lineno))

    '''

        @brief  remove text from image by using ocr image

        @author

    '''

    def removeText(self, img, pt, imgOCR):

        try:

            x = round(pt[0])

            y = round(pt[1])

            width, height = imgOCR.shape[::-1]

            imgOCR = cv2.dilate(imgOCR, np.ones((1,1), np.uint8))

            imgXOR = cv2.bitwise_xor(img[y:y+height, x:x+width], cv2.bitwise_not(imgOCR))

            imgXOR = cv2.dilate(imgXOR, np.ones((2,2), np.uint8))

            img[y:y+height, x:x+width] = imgXOR

        except Exception as ex:

            print('error occured({}) in {}:{}'.format(ex, sys.exc_info()[-1].tb_frame.f_code.co_filename, sys.exc_info()[-1].tb_lineno))

        return img