HYTOS

# coding: utf-8

import sys

import os, time, datetime

from PyQt5.QtCore import *

from PyQt5.QtGui import *

from PyQt5.QtWidgets import *

from AppDocData import *

import TrainingSymbolEditor_UI

import QtImageViewer

from GraphicsBoundingBoxItem import QGraphicsBoundingBoxItem

from TrainingBoxItem import QTrainingBoxItem

import cv2

import numpy as np

from PIL import Image

import AreaZoomCommand

from xml.etree.ElementTree import Element, SubElement, dump, ElementTree, parse

class QTrainingSymbolEditorDialog(QDialog):

    def __init__(self, parent, trainingImgPath, trainingLabelPath, boundaryLabelData, dateItem, boxItem):

        self.spinBoxFlag = False

        QDialog.__init__(self, parent)

        self.setWindowFlag(Qt.WindowMinMaxButtonsHint)

        self.trainingImgPath = trainingImgPath

        self.trainingLabelPath = trainingLabelPath

        self.boundaryLabelData = boundaryLabelData

        self.dateItem = dateItem

        self.boxItem = boxItem

        self.isChanged = False

        appDocData = AppDocData.instance()

        project = appDocData.getCurrentProject()

        self.ui = TrainingSymbolEditor_UI.Ui_TrainingSymbolEditorDialog()

        self.ui.setupUi(self)

        self.graphicsViewTrainingDrawing = QtImageViewer.QtImageViewer(self)

        self.graphicsViewTrainingDrawing.setParent(self.ui.centralWidget)

        self.graphicsViewTrainingDrawing.useDefaultCommand()

        self.ui.verticalLayoutTrainingDrawing.addWidget(self.graphicsViewTrainingDrawing)

        self.graphicsViewZoomDrawing = QGraphicsView(self)

        self.graphicsViewZoomDrawing.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOff)

        self.graphicsViewZoomDrawing.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)

        self.graphicsViewZoomDrawing.setParent(self.ui.leftSideWidget)

        self.ui.horizontalLayoutZoomDrawing.addWidget(self.graphicsViewZoomDrawing)

        # 학습 이미지 읽어서 메인 뷰에 그림, 사이드 뷰에 추가

        try:

            #trainingImgPath = os.path.join(project.getTrainingFilePath(), 'seed.seedF.exp0.tif')

            cvImg = cv2.cvtColor(cv2.imread(trainingImgPath), cv2.COLOR_BGR2GRAY)

            cvImg = cv2.threshold(cvImg, 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)[1]

            bytesPerLine = cvImg.shape[1]

            image = QImage(cvImg.data, cvImg.shape[1], cvImg.shape[0], bytesPerLine, QImage.Format_Indexed8)

            self.graphicsViewTrainingDrawing.setImage(image)

            # 사이드 뷰

            if type(image) is QPixmap:

                pixmap = image

            elif type(image) is QImage:

                pixmap = QPixmap.fromImage(image)

            scene = QGraphicsScene()

            self.graphicsViewZoomDrawing.setScene(scene)

            self._pixmapHandle = self.graphicsViewZoomDrawing.scene().addPixmap(pixmap)

            self.graphicsViewZoomDrawing.scene().setSceneRect(0, 0, pixmap.width(), pixmap.height())

            self.graphicsViewZoomDrawing.fitInView(self.graphicsViewZoomDrawing.scene().sceneRect(), Qt.KeepAspectRatioByExpanding)

            rect = QGraphicsRectItem(0, 0, 0, 0)

            pen = QPen(Qt.SolidLine)

            pen.setColor(Qt.green)

            pen.setWidthF(1)

            pen.setJoinStyle(Qt.MiterJoin)

            rect.setPen(pen)

            rect.setOpacity(0.7)

            self.graphicsViewZoomDrawing.scene().addItem(rect)

        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)

        # 박스 읽어서 메인 뷰에 그림

        try:

            boxList = []

            if self.boundaryLabelData is None: # already exist label file

                xml = parse(trainingLabelPath)

                root = xml.getroot()

                objects = list(root.iter('object'))

                for object_node in objects:

                    name = object_node.find('name').text

                    boxNode = object_node.find('bndbox')

                    xMin = int(boxNode.find('xmin').text)

                    yMin = int(boxNode.find('ymin').text)

                    xMax = int(boxNode.find('xmax').text)

                    yMax = int(boxNode.find('ymax').text)

                    boxList.append([xMin, yMin, xMax - xMin, yMax - yMin, name])

            else:

                for data in boundaryLabelData:

                    boxList.append([data[0], data[1], data[2], data[3], 'None'])

            for label in boxList:

                singleBox = QTrainingBoxItem(label[4], int(label[0]), int(label[1]), int(label[2]), int(label[3]))

                singleBox.transfer.onRemoved.connect(self.itemRemoved)

                singleBox.addTextItemToScene(self.ui, self.graphicsViewTrainingDrawing, self.graphicsViewZoomDrawing, self.spinBoxFlag)

        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)

        # label table

        self.ui.tableWidget.setSelectionMode(QAbstractItemView.SingleSelection) 

        self.ui.tableWidget.setColumnCount(1)

        self.ui.tableWidget.setSortingEnabled(True)

        self.ui.tableWidget.setHorizontalHeaderLabels([self.tr('Labels')])

        self.ui.tableWidget.horizontalHeaderItem(0).setSizeHint(QSize(30, 30))

        self.ui.tableWidget.horizontalHeader().setSectionResizeMode(0, QHeaderView.Stretch)

        self.ui.tableWidget.resizeColumnsToContents()

        self.ui.tableWidget.setEditTriggers(QAbstractItemView.NoEditTriggers)

        self.ui.tableWidget.verticalHeader().hide()

        self.ui.tableWidget.setRowCount(len(parent.labelList))

        row = 0

        for label in parent.labelList:

            textWidget = QTableWidgetItem(label[0])

            self.ui.tableWidget.setItem(row, 0, textWidget)

            row += 1

        self.ui.tableWidget.cellClicked.connect(self.listCellClicked)

        self.ui.pushButtonZoom.clicked.connect(self.onAreaZoom)

        self.ui.pushButtonSave.clicked.connect(self.pushButtonSaveClicked)

        self.ui.pushButtonCancel.clicked.connect(self.pushButtonCancelClicked)

        self.ui.pushButtonDelete.clicked.connect(self.pushButtonDeleteClicked)

        self.ui.pushButtonAdd.clicked.connect(self.pushButtonAddClicked)

        self.ui.pushButtonSplit.clicked.connect(self.pushButtonSplitClicked)

        self.ui.spinBoxLeft.valueChanged.connect(self.spinBoxChangedEvent)

        self.ui.spinBoxTop.valueChanged.connect(self.spinBoxChangedEvent)

        self.ui.spinBoxWidth.valueChanged.connect(self.spinBoxChangedEvent)

        self.ui.spinBoxHeight.valueChanged.connect(self.spinBoxChangedEvent)

        self.ui.lineEditChar.returnPressed.connect(self.pushButtonChangeClicked)

        self.removedItems = []

    def listCellClicked(self, row, col):

        self.isChanged = True

        items = self.graphicsViewTrainingDrawing.scene().selectedItems()

        label = self.ui.tableWidget.item(row, col).text()

        self.onCharChanged(label)

    '''

        @brief      split boxitem by button click

        @author     euisung

        @date       2018.10.31

    '''

    def pushButtonSplitClicked(self):

        self.isChanged = True

        items = self.graphicsViewTrainingDrawing.scene().selectedItems()

        if len(items) == 1:

            return

        rect = items[0].rect()

        secondBox = QTrainingBoxItem('', rect.x() + round(rect.width() / 2), rect.y(), round(rect.width() / 2), rect.height())

        secondBox.transfer.onRemoved.connect(self.itemRemoved)

        secondBox.addTextItemToScene(self.ui, self.graphicsViewTrainingDrawing, self.graphicsViewZoomDrawing, self.spinBoxFlag)

        items[0].setRect(rect.x(), rect.y(), round(rect.width() / 2), rect.height())

    '''

        @brief      add boxitem by button click

        @author     euisung

        @date       2018.10.17

    '''

    def pushButtonAddClicked(self):

        self.isChanged = True

        items = self.graphicsViewTrainingDrawing.scene().selectedItems()

        allItems = self.graphicsViewTrainingDrawing.scene().items()

        totalWidth = 0

        totalHeight = 0

        count = 0

        if len(items) == 1:

            rect = items[0].rect()

        elif len(allItems) > 1:

            for item in allItems:

                if type(item) is QGraphicsPixmapItem:

                    x = int(item.boundingRect().width() / 2)

                    y = int(item.boundingRect().height()/ 2)

                elif hasattr(item, 'rect'):

                    count += 1

                    totalWidth += item.rect().width()

                    totalHeight += item.rect().height()

            rect = QRectF(x, y, round(totalWidth / count), round(totalHeight / count))

        else:

            for item in allItems:

                if type(item) is QGraphicsPixmapItem:

                    x = int(item.boundingRect().width() / 2)

                    y = int(item.boundingRect().height()/ 2)

            rect = QRectF(x, y, 5, 10)

        singleBox = QTrainingBoxItem('', rect.x() + 3, rect.y(), rect.width(), rect.height())

        singleBox.transfer.onRemoved.connect(self.itemRemoved)

        singleBox.addTextItemToScene(self.ui, self.graphicsViewTrainingDrawing, self.graphicsViewZoomDrawing, self.spinBoxFlag)

    '''

        @brief      delete boxitem by button click

        @author     euisung

        @date       2018.10.17

    '''

    def pushButtonDeleteClicked(self):

        self.isChanged = True

        items = self.graphicsViewTrainingDrawing.scene().selectedItems()

        for item in items:

            item.transfer.onRemoved.emit(item)

    def onCharChanged(self, text):

        """

        change lineedit's text with given text

        """

        if text:

            self.ui.lineEditChar.setText(text)

            self.pushButtonChangeClicked()

    '''

        @brief      change boxitem char by button click

        @author     euisung

        @date       2018.10.17

    '''

    def pushButtonChangeClicked(self):

        self.isChanged = True

        items = self.graphicsViewTrainingDrawing.scene().selectedItems()

        label = self.ui.lineEditChar.text()

        if label == '' or label == 'None':

            return

        for box in [box for box in items if type(box) is QTrainingBoxItem]:

            box.char = label

            box.textShowBox.setHtml(label)

        rowCount = self.ui.tableWidget.rowCount()

        for index in range(rowCount):

            if self.ui.tableWidget.item(index, 0).text() == label:

                return

        self.ui.tableWidget.setRowCount(rowCount + 1)

        textWidget = QTableWidgetItem(label)

        self.ui.tableWidget.setItem(rowCount, 0, textWidget)

    '''

        @brief      close dialog by button click

        @author     euisung

        @date       2018.10.17

    '''

    def pushButtonCancelClicked(self):

        if self.isChanged:

            reply = QMessageBox.question(self, self.tr('Continue?'), self.tr('Changes may not have been saved.'), QMessageBox.Ignore, QMessageBox.Cancel)

            if reply == QMessageBox.Ignore:

                QDialog.reject(self)

        else:

            QDialog.reject(self)

    '''

        @brief      save box item by button click

        @author     euisung

        @date       2018.10.16

    '''

    def pushButtonSaveClicked(self):

        xml = Element('annotation')

        SubElement(xml, 'folder').text = 'None'

        SubElement(xml, 'filename').text = os.path.basename(self.trainingImgPath)

        pathNode = Element('path')

        pathNode.text = self.trainingImgPath.replace('/', '\\')

        xml.append(pathNode)

        sourceNode = Element('source')

        databaseNode = Element('database')

        databaseNode.text = 'Unknown'

        sourceNode.append(databaseNode)

        xml.append(sourceNode)

        items = self.graphicsViewTrainingDrawing.scene().items()

        for item in items:

            if type(item) is QGraphicsPixmapItem:

                back = item

                break

        sizeNode = Element('size')

        widthNode = Element('width')

        widthNode.text = str(int(back.boundingRect().width()))

        sizeNode.append(widthNode)

        heightNode = Element('height')

        heightNode.text = str(int(back.boundingRect().height()))

        sizeNode.append(heightNode)

        depthNode = Element('depth')

        depthNode.text = '3'

        sizeNode.append(depthNode)

        xml.append(sizeNode)

        segmentedNode = Element('segmented')

        segmentedNode.text = '0'

        xml.append(segmentedNode)

        labelContent = []

        counts = {}

        for item in items:

            if type(item) is QTrainingBoxItem:

                rect = item.rect()

                label, xMin, yMin, xMax, yMax = item.char, int(rect.x()), int(rect.y()), int(rect.x()) + int(rect.width()), int(rect.y()) + int(rect.height())

                if label == 'None' or label == '':

                    continue

                if label not in labelContent:

                    labelContent.append(label)

                    counts[label] = 1

                else:

                    counts[label] = counts[label] + 1

                objectNode = Element('object')

                nameNode = Element('name')

                nameNode.text = label

                objectNode.append(nameNode)

                poseNode = Element('pose')

                poseNode.text = 'Unspecified'

                objectNode.append(poseNode)

                truncatedNode = Element('truncated')

                truncatedNode.text = '0'

                objectNode.append(truncatedNode)

                difficultNode = Element('difficult')

                difficultNode.text = '0'

                objectNode.append(difficultNode)

                bndboxNode = Element('bndbox')

                xminNode = Element('xmin')

                xminNode.text = str(xMin)

                bndboxNode.append(xminNode)

                yminNode = Element('ymin')

                yminNode.text = str(yMin)

                bndboxNode.append(yminNode)

                xmaxNode = Element('xmax')

                xmaxNode.text = str(xMax)

                bndboxNode.append(xmaxNode)

                ymaxNode = Element('ymax')

                ymaxNode.text = str(yMax)

                bndboxNode.append(ymaxNode)

                objectNode.append(bndboxNode)

                xml.append(objectNode)

        ElementTree(xml).write(self.trainingLabelPath)

        modifiedTime = str(datetime.datetime.strptime(time.ctime(os.path.getmtime(self.trainingLabelPath)), "%a %b %d %H:%M:%S %Y"))

        self.dateItem.setText(modifiedTime)

        labelCount = [label + ':' + str(counts[label]) for label in labelContent]

        labels = ','.join(labelCount)

        self.boxItem.setText(labels)

        self.isChanged = False

        reply = QMessageBox.information(self, self.tr('Information'), self.tr('Save Success!'), QMessageBox.Ok, QMessageBox.Cancel)

        if reply == QMessageBox.Ok:

            QDialog.reject(self)

    '''

        @brief      chane box item's rect by button click

        @author     euisung

        @date       2018.10.16

    '''

    def spinBoxChangedEvent(self, event):

        self.isChanged = True

        items = self.graphicsViewTrainingDrawing.scene().selectedItems()

        if (len(items) != 1) or self.spinBoxFlag:

            return

        spinBoxName = self.sender().objectName()

        rect = items[0].rect()

        bound = self.graphicsViewZoomDrawing.scene().items()[0]

        if spinBoxName == 'spinBoxLeft':

            spinBoxValue = self.ui.spinBoxLeft.value()

            items[0].setRect(QRectF(spinBoxValue, rect.y(), rect.width(), rect.height()))

            items[0].setPosCustom()

        elif spinBoxName == 'spinBoxTop':

            spinBoxValue = self.ui.spinBoxTop.value()

            items[0].setRect(QRectF(rect.x(), spinBoxValue, rect.width(), rect.height()))

            items[0].setPosCustom()

        elif spinBoxName == 'spinBoxWidth':

            spinBoxValue = self.ui.spinBoxWidth.value()

            items[0].setRect(QRectF(rect.x(),  rect.y(), spinBoxValue, rect.height()))

            items[0].setPosCustom()

        elif spinBoxName == 'spinBoxHeight':

            spinBoxValue = self.ui.spinBoxHeight.value()

            items[0].setRect(QRectF(rect.x(), rect.y(), rect.width(), spinBoxValue))

            items[0].setPosCustom()

        rect = items[0].rect()        

        bound.setRect(rect)

        rectSide = QRectF(rect.x() - 3, rect.y() - 3, rect.width() + 6,  rect.height() + 6)

        self.graphicsViewZoomDrawing.fitInView(rectSide)

        self.graphicsViewTrainingDrawing.scene().update()

    def onAreaZoom(self, action):

        if self.ui.pushButtonZoom.isChecked():

            cmd = AreaZoomCommand.AreaZoomCommand(self.graphicsViewTrainingDrawing)

            cmd.onRejected.connect(self.onCommandRejected)

            self.graphicsViewTrainingDrawing.command = cmd

    def onCommandRejected(self, cmd):

        try:

            if type(cmd) is AreaZoomCommand.AreaZoomCommand:

                self.ui.pushButtonZoom.setChecked(False)

        finally:

            self.graphicsViewTrainingDrawing.useDefaultCommand()

    def itemRemoved(self, item):

        try:

            if type(item) is QTrainingBoxItem:

                self.removedItems.append(str(item.uid))

            if item.scene is not None: item.scene().removeItem(item)

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

    '''

        @brief      key pressed event

        @author     euisung

        @date       2018.11.05

    '''

    def keyPressEvent(self, event):

        try:

            items = self.graphicsViewTrainingDrawing.scene().selectedItems()

            if items and type(items[0]) is not QTrainingBoxItem:

                return

            selectedItem = items[0]

            closestItem = None

            x = selectedItem.rect().x()

            if event.key() == Qt.Key_Left:

                dx =  sys.maxsize

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

                    if type(item) is QTrainingBoxItem:

                        if x - item.rect().x() > 0 and x - item.rect().x() < dx:

                            closestItem = item

                            dx = x - item.rect().x()

                if closestItem is not None:

                    closestItem.setSelected(True)

                    selectedItem.setSelected(False)

                    closestItem.mousePressEvent('arrow key')             

            elif event.key() == Qt.Key_Right:

                dx = -sys.maxsize

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

                    if type(item) is QTrainingBoxItem:

                        if x - item.rect().x() < 0 and x - item.rect().x() > dx:

                            closestItem = item

                            dx = x - item.rect().x()

                if closestItem is not None:

                    closestItem.setSelected(True)

                    selectedItem.setSelected(False)

                    closestItem.mousePressEvent('arrow key')

            elif event.key() == Qt.Key_Up:

                self.ui.spinBoxWidth.setValue(self.ui.spinBoxWidth.value() + 1)

            elif event.key() == Qt.Key_Down:

                self.ui.spinBoxWidth.setValue(self.ui.spinBoxWidth.value() - 1)

            elif event.key() == Qt.Key_Return and not self.ui.lineEditChar.hasFocus():

                if self.ui.tableWidget.selectedItems():

                    self.isChanged = True

                    self.onCharChanged(self.ui.tableWidget.selectedItems()[0].text())

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

            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)

            return None