HYTOS

# coding: utf-8

"""

    This is ocr result dialog module

"""

from PIL import Image

import io

import numpy as np

import math

import enum

from PyQt5.QtCore import *

from PyQt5.QtGui import *

from PyQt5.QtWidgets import *

import OcrResultDialog_UI

import QtImageViewer

import tesseract_ocr_module as TOCR

from App import App

from AppDocData import *

from TextInfo import TextInfo

class SpellTextEdit(QTextEdit):

    def __init__(self, *args):

        QTextEdit.__init__(self, *args)

        # Default dictionary based on the current locale.

        app_doc_data = AppDocData.instance()

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

        self.highlighter = Highlighter(self.document())

        self.highlighter.white_char_list = white_char_list[0].value if white_char_list else None

class Highlighter(QSyntaxHighlighter):

    err_format = QTextCharFormat()

    err_format.setUnderlineColor(Qt.red)

    err_format.setUnderlineStyle(QTextCharFormat.SpellCheckUnderline)

    def __init__(self, *args):

        QSyntaxHighlighter.__init__(self, *args)

        self.white_char_list = None

    def highlightBlock(self, text):

        pos = 0

        for word in text.split():

            if self.white_char_list and any((c not in self.white_char_list) for c in word):

                self.setFormat(pos, len(word), self.err_format)

            pos += len(word) + 1

class QOcrResultDialog(QDialog):

    class Format(enum.Enum):

        Normal = 0

        Table = 1

    def __init__(self, parent, qimage, boundingBox, format=Format.Normal, text_item=None):

        QDialog.__init__(self, parent)

        self.textInfoList = []

        self._text_item = text_item

        self.image = qimage

        self.boundingBox = boundingBox

        self._format = format

        self.angle = 0  # angle in degree

        self.ui = OcrResultDialog_UI.Ui_Dialog()

        self.ui.setupUi(self)

        self.ui.detectResultTextEdit = SpellTextEdit()

        self.ui.detectResultTextEdit.setFont(QFont('Consolas', 15, QFont.Bold))

        self.ui.horizontalLayoutTextEdit.addWidget(self.ui.detectResultTextEdit)

        app_doc_data = AppDocData.instance()

        configs = app_doc_data.getAppConfigs('app', 'mode')

        if configs and 1 == len(configs) and 'advanced' == configs[0].value:

            pass

        else:

            self.ui.pushButtonMakeTrainingImage.setVisible(False)

        self.graphicsView = QtImageViewer.QtImageViewer(App.mainWnd())

        self.graphicsView.useDefaultCommand()  # USE DEFAULT COMMAND

        self.graphicsView.setImage(self.image)

        self.ui.horizontalLayoutGraphicsView.addWidget(self.graphicsView)

        self.ui.counterClockPushButton_2.clicked.connect(lambda: self.rotateImage(True))

        self.ui.clockPushButton_2.clicked.connect(lambda: self.rotateImage(False))

        self.ui.pushButtonCopyHori.clicked.connect(self.copy_horizontal)

        # add shortcut for detecting text with 't'

        self.ui.redetectPushButton.clicked.connect(self.detect_text)

        shortcut = QShortcut(QKeySequence('t'), self.ui.redetectPushButton)

        shortcut.activated.connect(self.detect_text)

        # up to here

        self.ui.pushButtonMakeTrainingImage.clicked.connect(self.pushButtonMakeTrainingImageClicked)

        self.ui.comboBoxOCRData.addItem('eng')

        tessdata_path = os.path.join(os.getenv('ALLUSERSPROFILE'), 'Digital PID', 'Tesseract-OCR', 'tessdata')

        if os.path.isfile(os.path.join(tessdata_path, app_doc_data.getCurrentProject().name + '.traineddata')):

            self.ui.comboBoxOCRData.addItem(app_doc_data.getCurrentProject().name)

        configs = app_doc_data.getConfigs('Text Recognition', 'OCR Data')

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

        if value:

            at = self.ui.comboBoxOCRData.findText(value)

            self.ui.comboBoxOCRData.setCurrentIndex(at)

        else:

            self.ui.comboBoxOCRData.selectedIndex = 0

        if not self._text_item:

            if format == QOcrResultDialog.Format.Normal and self.boundingBox.height() > self.boundingBox.width():

                self.rotateImage(False)

            self.detect_text()

        else:

            rect = self._text_item.sceneBoundingRect()

            text_info = TextInfo(self._text_item.text(), 0, 0, rect.width(), rect.height(), 0)

            self.textInfoList.append(text_info)

            self.display_text_rect()

            allowed_error = 0.001

            if abs(self._text_item.angle - 1.57) < allowed_error or abs(self._text_item.angle - 4.71) < allowed_error:

                self.rotateImage(False)

            self.ui.detectResultTextEdit.setPlainText(self._text_item.text())

            self.ui.checkBoxSeperate.setChecked(False)

        self.isAccepted = False

    def showEvent(self, QShowEvent):

        """show event"""

        self.graphicsView.zoomImageInit()

    def display_text_rect(self):

        """display text bounding rectangle"""

        for item in self.graphicsView.scene.items():

            if type(item) is not QGraphicsPixmapItem:

                self.graphicsView.scene.removeItem(item)

        for text_info in self.textInfoList:

            self.graphicsView.scene.addRect(text_info.getX(), text_info.getY(),

                                            text_info.getW(), text_info.getH(), QPen(Qt.red, 1, Qt.SolidLine))

    '''

        @brief      Make OCR Training Image

        @author     euisung

        @date       2018.10.16

        @history    euisung     2018.11.02       add notice push

    '''

    def pushButtonMakeTrainingImageClicked(self):

        import uuid

        uid = str(uuid.uuid4()) + '.png'

        appDocData = AppDocData.instance()

        project = appDocData.getCurrentProject()

        trainingImgPath = os.path.join(project.getTrainingFilePath(), uid)

        self.image.save(trainingImgPath)

        QMessageBox.about(self, self.tr("INFO"), self.tr('Successfully saved.'))

        QDialog.reject(self)

    def rotateImage(self, isCounterClock):

        """rotate the image"""

        transform = QTransform()

        if isCounterClock:

            '''CounterClock'''

            self.angle = (self.angle - 90) % 360

            transform.rotate(-90)

        else:

            '''Clock'''

            self.angle = (self.angle - 270) % 360

            transform.rotate(90)

        self.graphicsView.clearImage()

        self.image = self.image.transformed(transform)

        self.graphicsView.setImage(self.image)

        for text_info in self.textInfoList:

            rect = QRectF(text_info.getX(), text_info.getY(), text_info.getW(), text_info.getH())

            rect = transform.mapRect(rect)

            text_info.setX(self.image.width() + rect.left() if rect.left() < 0 else rect.left())

            text_info.setY(self.image.height() - max(abs(rect.top()), abs(rect.bottom())) if rect.top() < 0 else rect.top())

            text_info.setW(rect.width())

            text_info.setH(rect.height())

            self.graphicsView.scene.addRect(QRectF(text_info.getX(), text_info.getY(), text_info.getW(), text_info.getH()),

                                            QPen(Qt.red, 1, Qt.SolidLine))

    '''

        @history 2018.04.26 Jeongwoo    Add Rectangle with modified Coords

                 2018.06.20 Jeongwoo    Remove test code

                 2018.11.08 euisung     add white char list check process on db

                 2018.11.22 euisung     OCR lang apply fixed

    '''

    def detect_text(self):

        import cv2

        from TextDetector import TextDetector

        from TextInfo import TextInfo

        try:

            buffer = QBuffer()

            buffer.open(QBuffer.ReadWrite)

            self.image.save(buffer, "PNG")

            pyImage = Image.open(io.BytesIO(buffer.data()))

            img = np.array(pyImage)

            if len(img.shape[::-1]) == 2:

                img_width, img_height = img.shape[::-1]

            else:

                _, img_width, img_height = img.shape[::-1]

            app_doc_data = AppDocData.instance()

            ocr_data = self.ui.comboBoxOCRData.currentText()

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

            if self._format == QOcrResultDialog.Format.Normal:

                self.textInfoList = TOCR.getTextInfo(img, (0, 0), 0, language=ocr_data,

                                                     conf=white_char_list[0].value if white_char_list else '')

                if self.textInfoList:

                    self.ui.detectResultTextEdit.setText(self.getPlainText(self.textInfoList))

                    self.display_text_rect()

                    self.copy_horizontal()

                else:

                    self.ui.detectResultTextEdit.setText(self.tr("Not Found"))

            else:

                cv_image = cv2.threshold(cv2.cvtColor(img, cv2.COLOR_BGR2GRAY), 200, 255, cv2.THRESH_BINARY)[1]

                # contours 추출

                contours, _ = cv2.findContours(cv_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

                text_info_list = []

                for contour in contours:

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

                    text_info_list.append(TextInfo('', x, y, w, h, 0))

                text_info_list = [text_info for text_info in text_info_list if

                                  not any([_text_info for _text_info in text_info_list if _text_info is not text_info

                                           and text_info.contains(_text_info.center)])]

                detector = TextDetector()

                detector.recognizeText(img, (0, 0), text_info_list, None, None, None, None, onlyTextArea=True)

                self.textInfoList = detector.textInfoList.copy()

                self.textInfoList.sort(key=lambda x: x.getY())

                """

                index = 0

                for rect in rects:

                    cropped = img[rect.y():rect.y() + rect.height(), rect.x():rect.x() + rect.width()]

                    cv2.imwrite(f"c:\\temp\\ocr-{index}.png", cropped)

                    text_info = TOCR.getTextInfo(cropped, (0, 0), 0, language=ocr_data,

                                                 conf=white_char_list[0].value if white_char_list else '')

                    if text_info:

                        for _text_info in text_info:

                            _text_info.setX(_text_info.getX() + rect.x())

                            _text_info.setY(_text_info.getY() + rect.y())

                        self.textInfoList.extend(text_info)

                    index = index + 1

                """

                if self.textInfoList:

                    self.ui.detectResultTextEdit.setText(self.getPlainText(self.textInfoList))

                    self.display_text_rect()

                    self.copy_horizontal()

                else:

                    self.ui.detectResultTextEdit.setText(self.tr("Not Found"))

        except Exception as ex:

            from App import App

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

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

            App.mainWnd().addMessage.emit(MessageType.Error, message)

    def getPlainText(self, textInfoList):

        text = ''

        for index in range(len(textInfoList)):

            textInfo = textInfoList[index]

            if index != 0:

                text = text + '\n'

            text = text + textInfo.getText()

        return text

    '''

        @brief      OK Button Clicked. Remake TextInfo object

        @author     Jeongwoo

        @date       18.04.19

        @history    18.04.20    Jeongwoo    Calculate Start Point Coordinates by rotated angle

                    18.04.26    Jeongwoo    Scene.itemAt(textX - boundBox.x(), textY - boundBox.y())

    '''

    def accept(self):

        from TextInfo import TextInfo

        self.isAccepted = True

        try:

            text = self.ui.detectResultTextEdit.toPlainText()

            if text == '' or text == 'Not Found':

                QMessageBox.about(self.ui.ocrDialogButtonBox, self.tr('Notice'),

                                  self.tr('Please try again after recognition or type.'))

                return

            isSplit = self.ui.checkBoxSeperate.isChecked()

            if isSplit:

                splitText = text.split('\n')

            else:

                splitText = [text]

            # try to detect text if there is no result of detection or

            # count of text info list not match with count of split text

            if isSplit:

                if self.textInfoList and (len(self.textInfoList) == len(splitText)):

                    for index in range(len(self.textInfoList)):

                        self.textInfoList[index].setText(splitText[index])

                elif not self.textInfoList:

                    self.detect_text()

                    if len(self.textInfoList) == len(splitText):

                        for index in range(len(self.textInfoList)):

                            self.textInfoList[index].setText(splitText[index])

                    else:

                        self.textInfoList = self.getMergedTextInfo(text)

            elif len(self.textInfoList) > 1 or len(self.textInfoList) == 0:

                self.textInfoList = self.getMergedTextInfo(text)

            radian = round(math.radians(abs(self.angle)), 2)

            for idx in range(len(self.textInfoList)):

                text_info = self.textInfoList[idx]

                # update text using user input text

                if idx < len(splitText):

                    text_info.setText(splitText[idx])

                # up to here

                if radian == 1.57 or radian == 4.71:

                    text_info.setAngle(radian)  # 360 degree == 6.28319 radian

                    # rotate text information

                    trans = QTransform()

                    trans.rotate(self.angle*-1)

                    rect = QRectF(text_info.getX(), text_info.getY(), text_info.getW(), text_info.getH())

                    rect = trans.mapRect(rect)

                    width, height = self.image.height(), self.image.width()

                    x = width + rect.left() if rect.left() < 0 else rect.left()

                    y = height - max(abs(rect.top()), abs(rect.bottom())) if rect.top() < 0 else rect.top()

                    text_info.setX(self.boundingBox.x() + x)

                    text_info.setY(self.boundingBox.y() + y)

                    text_info.setW(rect.width())

                    text_info.setH(rect.height())

                    # up to here

                else:

                    text_info.setX(int(self.boundingBox.x()) + text_info.getX())

                    text_info.setY(int(self.boundingBox.y()) + text_info.getY())

            QDialog.accept(self)

        except Exception as ex:

            from App import App

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

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

            App.mainWnd().addMessage.emit(MessageType.Error, message)

    def getMergedTextInfo(self, text):

        import cv2

        from TextInfo import TextInfo

        buffer = QBuffer()

        buffer.open(QBuffer.ReadWrite)

        self.image.save(buffer, "PNG")

        pyImage = Image.open(io.BytesIO(buffer.data()))

        img = np.array(pyImage)

        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        imgNot = np.ones(img.shape, np.uint8)

        cv2.bitwise_not(img, imgNot)

        imgNot = cv2.dilate(imgNot, np.ones((8, 8), np.uint8))

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

        minX, minY, maxX, maxY = sys.maxsize, sys.maxsize, 0, 0

        if len(contours) is 0:

            minX, minY, maxX, maxY = 0, 0, self.image.width(), self.image.height()

        else:

            minX, minY, maxX, maxY = sys.maxsize, sys.maxsize, 0, 0

            for cnt in contours:

                x, y, w, h = cv2.boundingRect(cnt)

                minX = min(x, minX)

                minY = min(y, minY)

                maxX = max(x + w, maxX)

                maxY = max(y + h, maxY)

            minX, minY, maxX, maxY = minX, minY, maxX, maxY

        return [TextInfo(text, minX, minY, maxX - minX, maxY - minY, 0)]

    def reject(self):

        self.isAccepted = False

        self.textInfoList = None

        QDialog.reject(self)

    def copy_horizontal(self):

        import io, csv

        try:

            table = [[text for text in self.ui.detectResultTextEdit.toPlainText().split('\n')]]

            stream = io.StringIO()

            csv.writer(stream, delimiter='\t').writerows(table)

            QApplication.clipboard().setText(stream.getvalue())

        except Exception as ex:

            from App import App 

            from AppDocData import MessageType

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

            App.mainWnd().addMessage.emit(MessageType.Error, message)

    '''

        @brief  Display this QDialog

    '''

    def showDialog(self):

        # self.setWindowFlags(self.windowFlags() & ~Qt.WindowContextHelpButtonHint)

        res = self.exec_()

        return res, self.textInfoList