HYTOS

# coding: utf-8

"""

    This is ocr result dialog module

"""

from PIL import Image

import io, sys

import numpy as np

import math

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 *

class QOcrResultDialog(QDialog):

    def __init__(self, parent, qimage, boundingBox, isModify = False, text = None):

        QDialog.__init__(self, parent)

        self.textInfoList = []

        self.isModify = isModify

        self.image = qimage

        self.originImageWidth = qimage.width()

        self.originImageHeight = qimage.height()

        self.boundingBox = boundingBox

        self.angle = 0 # angle is degree

        self.ui = OcrResultDialog_UI.Ui_Dialog()

        self.ui.setupUi(self)

        appDocData = AppDocData.instance()

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

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

            pass

        else:

            self.ui.pushButtonMakeTrainingImage.setVisible(False)

        self.imgW = qimage.width()

        self.imgH = qimage.height()

        self.image = self.image.scaled(self.imgW, self.imgH)

        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.redetectPushButton_2.clicked.connect(self.detectText)

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

        if self.isModify == False:

            self.detectText()

        else:

            self.ui.detectResultTextEdit.setPlainText(text)

        self.isAccepted = False

    '''

        @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):

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

            self.graphicsView.scene.removeItem(item)

        self.graphicsView.clearImage()

        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)

        #print(str(360 - self.angle))

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

        self.graphicsView.setImage(self.image)

        self.textInfoList = []

    '''

        @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 detectText(self):

        try:

            buffer = QBuffer()

            buffer.open(QBuffer.ReadWrite)

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

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

            img = np.array(pyImage)

            #self.image.save('c:\\temp\\a.png')

            docData = AppDocData.instance()

            # get ocr data of area which has the text

            pt = self.boundingBox.center()

            areas = [area for area in docData.getAreaList() if area.contains((pt.x(), pt.y()))]

            ocr_data = sorted(areas, key=lambda attr: attr.area)[0].OCRData if areas else 'eng'

            # up to here

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

            if len(whiteCharList) is 0:

                self.textInfoList = TOCR.getTextInfo(img, (round(self.boundingBox.x()), round(self.boundingBox.y())), 0, language=ocr_data)

            else:

                self.textInfoList = TOCR.getTextInfo(img, (round(self.boundingBox.x()), round(self.boundingBox.y())), 0, language=ocr_data, conf = whiteCharList[0].value)

            if self.textInfoList is not None and len(self.textInfoList) > 0:

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

                for textInfo in self.textInfoList:

                    self.graphicsView.scene.addRect(textInfo.getX()-int(self.boundingBox.x()), textInfo.getY()-int(self.boundingBox.y()), textInfo.getW(), textInfo.getH(), QPen(Qt.red, 1, Qt.SolidLine))

            else:

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

        except Exception as ex:

            from App import App 

            message = 'error occured({}) 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 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, 'Notice', 'Please try again after recognition or type.')

                return

            isSplit = self.ui.checkBoxSeperate.isChecked()

            if isSplit:

                splitText = text.split('\n')

            else:

                splitText = [text]

            if not len(self.textInfoList) > 0:

                import cv2

                #QMessageBox.about(self.ui.ocrDialogButtonBox, "알림", "텍스트 검출을 하신 후 다시 시도해주세요.")

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

                image, 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 = self.boundingBox.x(), self.boundingBox.y(), self.boundingBox.x() + self.image.width(), self.boundingBox.y() + 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 + self.boundingBox.x(), minY + self.boundingBox.y(), maxX + self.boundingBox.x(), maxY + self.boundingBox.y()

                self.textInfoList.append(TextInfo(text, minX, minY, maxX - minX, maxY - minY, 0))

            if not isSplit:

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

                for textInfo in self.textInfoList:

                    x, y, w, h = textInfo.getX() ,textInfo.getY(), textInfo.getW(), textInfo.getH()

                    minX = min(x ,minX)

                    minY = min(y, minY)

                    maxX = max(x + w, maxX)

                    maxY = max(y + h, maxY)

                self.textInfoList = [TextInfo(text, minX, minY, maxX - minX, maxY - minY, 0)]

            if len(self.textInfoList) > 0:

                for index in range(len(splitText)):

                    textInfo = self.textInfoList[index]

                    item = self.graphicsView.scene.itemAt(QPointF(float(textInfo.getX() - int(self.boundingBox.x())), float(textInfo.getY()-int(self.boundingBox.y()))), QTransform())

                    if item is not None:

                        ## Transform rectangle for calculate start point

                        imgTransform = QTransform()

                        if self.angle == 90 or self.angle == 270:

                            imgTransform.translate(self.image.height()*0.5, self.image.width()*0.5)

                        elif self.angle == 0 or self.angle == 360:

                            imgTransform.translate(self.image.width()*0.5, self.image.height()*0.5)

                        imgTransform.rotate(-abs(self.angle))

                        imgTransform.translate(-self.image.width()*0.5, -self.image.height()*0.5)

                        rect = QRect(textInfo.getX() - int(self.boundingBox.x()), textInfo.getY() - int(self.boundingBox.y()), textInfo.getW(), textInfo.getH())

                        rect = imgTransform.mapRect(rect)

                        ## up to here

                        textInfo.setX(rect.x() + int(self.boundingBox.x()))

                        textInfo.setY(rect.y() + int(self.boundingBox.y()))

                        textInfo.setText(splitText[index])

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

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

                        if radian == 1.57 or radian == 4.71:

                            width = textInfo.getW()

                            height = textInfo.getH()

                            textInfo.setW(height) ## SWAP

                            textInfo.setH(width) ## SWAP

                self.textInfoList = self.textInfoList[:len(splitText)]

                QDialog.accept(self)

        except Exception as ex:

            from App import App 

            message = 'error occured({}) 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 reject(self):

        self.isAccepted = False

        self.textInfoList = None

        QDialog.reject(self)

    '''

        @brief  Display this QDialog

    '''

    def showDialog(self):

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

        self.exec_()

        return (self.isAccepted, self.textInfoList)