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 = []

        self.titleBlockTextInfoList = []

    '''

        @brief  detect text areas

        @author humkyung

        @date   2018.06.16

    '''

    def detectTextAreas(self, img, offset):

        try:

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

        except Exception as ex:

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

                                                       sys.exc_info()[-1].tb_lineno))

        return None, None

    '''

        @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

        list = []

        ocr_image = None

        try:

            app_doc_data = AppDocData.instance()

            project = app_doc_data.getCurrentProject()

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

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

            minSize = 5

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

            # binaryImg, mask = cv2.threshold(imgGray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

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

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

            for contour in contours:

                # remove too big one or horizontal/vertical line

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

                area = cv2.contourArea(contour, True)

                # skip one which size is greater than max size or less then minimum size

                if area >= 0:

                    if (w > maxTextSize or h > maxTextSize) or (w <= minSize and h <= minSize):

                        continue

                if area >= 0:

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

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

                else:

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

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

            cv2.imwrite(path, ocr_image)

            rects = []

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

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

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

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

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

            eroded = cv2.bitwise_not(eroded)

            bboxes = []

            contours, hierarchy = cv2.findContours(eroded, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

            for contour in contours:

                area = cv2.contourArea(contour, True)

                if area < 0:

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

                    bboxes.append(QRect(x, y, w, h))

            # exclude bounding boxes contains child bounding box

            not_containing_bbox = []

            for bbox in bboxes:

                matches = [_bbox for _bbox in bboxes if bbox != _bbox and bbox.contains(_bbox)]

                if not matches:

                    not_containing_bbox.append(bbox)

            # up to here

            for bbox in not_containing_bbox:

                x, y = bbox.left(), bbox.top()

                w, h = bbox.width(), bbox.height()

                img = ocr_image[bbox.top():bbox.bottom(), bbox.left():bbox.right()]

                img = cv2.bitwise_not(img)

                horizontal, max_width = 0, 0

                vertical, max_height = 0, 0

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

                for xx in _contours:

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

                    max_width = _x if _x > max_width else max_width

                    max_height = _y if _y > max_height else max_height

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

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

                    else:

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

                if (w < minSize and h < minSize) or (max_width > maxTextSize and max_height > maxTextSize):

                    continue  # skip too small or big one

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

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

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

            # merge rectangles

            interestings = []

            while rects:

                rect = rects.pop()

                if 0 == rect[0]:    # x-direction text

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

                    matches = [x for x in rects if (x[0] == rect[0]) and

                               abs(x[1].height() - rect[1].height()) < (x[1].height() + rect[1].height())*0.5 and

                               abs(x[1].center().y() - rect[1].center().y()) < rect[1].height()*0.5 and

                               rectExpand.intersects(x[1])]

                else:               # y -direction text

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

                    matches = [x for x in rects if (x[0] == rect[0]) and

                               abs(x[1].width() - rect[1].width()) < (x[1].width() + rect[1].width())*0.5 and

                               abs(x[1].center().x() - rect[1].center().x()) < rect[1].width()*0.5 and

                               rectExpand.intersects(x[1])]

                if matches:

                    for _rect in matches:

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

                        if _rect in rects:

                            rects.remove(_rect)

                    rects.append(rect)

                else:

                    interestings.append(rect)

            for rect in interestings:

                matches = [_rect for _rect in interestings if rect != _rect and _rect[1].contains(rect[1])]

                # if there is no boxes which contains

                if not matches:

                    angle = rect[0]

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

                                            rect[1].height(), angle))

        except Exception as ex:

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

                                                           sys.exc_info()[-1].tb_lineno)

            print(message)

        return list, ocr_image

    '''

        @brief      recognize text of given text info

        @author     humkyung

        @date       2018.07.24

        @history    change parameter updateProgressSignal to worker

                    2018.11.08 euisung     add white char list check process on db

    '''

    @staticmethod

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

        import re

        res = []

        app_doc_data = AppDocData.instance()

        try:

            for tInfo in tInfos:

                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 = app_doc_data.getSymbolCategoryByType(_type)

                    if 'Instrumentation' == category:

                        tInfo.setAngle(0)

                # up to here

                whiteCharList = app_doc_data.getConfigs('Text Recognition', 'White Character List')

                if len(whiteCharList) is 0:

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

                else:

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

                                                      conf=whiteCharList[0].value)

                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() if resultTextInfo is not None and 1 == len(

                                                   resultTextInfo) else None)

        except Exception as ex:

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

                                                           sys.exc_info()[-1].tb_lineno)

            if worker is not None:

                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

                    euisung  2018.11.08 add white char list check process on db

                    euisung  2018.11.12 add title block properties

    '''

    def recognizeText(self, imgSrc, offset, tInfoList, searchedSymbolList, worker, listWidget, maxProgressValue,

                      onlyTextArea=False):

        import concurrent.futures as futures

        from multiprocessing import Process, Queue

        from Area import Area

        try:

            self.otherTextInfoList = []

            self.titleBlockTextInfoList = []

            self.textInfoList = []

            app_doc_data = AppDocData.instance()

            project = app_doc_data.getCurrentProject()

            text_info_array = np.array_split(tInfoList, THREAD_MAX_WORKER if len(tInfoList) > THREAD_MAX_WORKER else \

                len(tInfoList))

            with futures.ThreadPoolExecutor(max_workers=THREAD_MAX_WORKER) as pool:

                future_text = {pool.submit(TextDetector.recognizeTextFromImage, tInfo, imgSrc, offset,

                                       searchedSymbolList, worker, listWidget, maxProgressValue):

                               tInfo for tInfo in text_info_array}

                for future in futures.as_completed(future_text):

                    try:

                        data = future.result()

                        if data:

                            self.textInfoList.extend(data)

                    except Exception as ex:

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

                                                                       sys.exc_info()[-1].tb_lineno)

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

            if onlyTextArea:

                return

            # parse texts in area except Drawing area

            whiteCharList = app_doc_data.getConfigs('Text Recognition', 'White Character List')

            for area in app_doc_data.getAreaList():

                if area.name == 'Drawing': continue

                if area.name == 'Note':

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

                        if len(whiteCharList) is 0:

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

                        else:

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

                                                     conf=whiteCharList[0].value)

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

                else:

                    img = app_doc_data.imgSrc[round(area.y):round(area.y + area.height),

                          round(area.x):round(area.x + area.width)]

                    if len(whiteCharList) is 0:

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

                    else:

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

                                                 conf=whiteCharList[0].value)

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

                        if area.name == 'Unit':

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

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

            titleBlockProps = app_doc_data.getTitleBlockProperties()

            if titleBlockProps:

                for titleBlockProp in titleBlockProps:

                    area = Area(titleBlockProp[0])

                    area.parse(titleBlockProp[2])

                    if not (titleBlockProp[3] and titleBlockProp[3] != ''):

                        img = app_doc_data.imgSrc[round(area.y):round(area.y + area.height),

                              round(area.x):round(area.x + area.width)]

                        if len(whiteCharList) is 0:

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

                        else:

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

                                                     conf=whiteCharList[0].value)

                        texts = [ti.TextInfo('\n'.join([textInfo.getText() for textInfo in texts]), area.x, area.y,

                                             area.width, area.height, 0)]

                    else:

                        texts = [ti.TextInfo(titleBlockProp[3], area.x, area.y, area.width, area.height, 0)]

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

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

            """

            for text_box in tInfoList:

                x = text_box.getX()

                y = text_box.getY()

                cv2.rectangle(imgSrc, (x - offset[0], y - offset[1]),

                              (x - offset[0] + text_box.getW(), y - offset[1] + text_box.getH()), 1, 1)

            cv2.imwrite('c:\\Temp\\text_box.png', imgSrc)

            """

        except Exception as ex:

            message = 'error occurred({}) in {}:{}'.format(repr(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 remove_text_from_image(self, imgSrc, offset):

        # 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())

            cv2.rectangle(imgSrc, (x, y), (x + width, y + height), 255, -1)

        # up to here

        # DEBUG

        #cv2.imwrite("c:\\temp\\remove_texts.png", imgSrc)