pull/7044/head
whjdark 2022-07-29 16:33:41 +08:00
parent 1696b36bdb
commit 97fc98cdd3
1 changed files with 74 additions and 3 deletions

View File

@ -21,6 +21,7 @@ import os.path
import platform import platform
import subprocess import subprocess
import sys import sys
from cv2 import INTER_MAX
import xlrd import xlrd
from functools import partial from functools import partial
@ -28,7 +29,7 @@ from PyQt5.QtCore import QSize, Qt, QPoint, QByteArray, QTimer, QFileInfo, QPoin
from PyQt5.QtGui import QImage, QCursor, QPixmap, QImageReader from PyQt5.QtGui import QImage, QCursor, QPixmap, QImageReader
from PyQt5.QtWidgets import QMainWindow, QListWidget, QVBoxLayout, QToolButton, QHBoxLayout, QDockWidget, QWidget, \ from PyQt5.QtWidgets import QMainWindow, QListWidget, QVBoxLayout, QToolButton, QHBoxLayout, QDockWidget, QWidget, \
QSlider, QGraphicsOpacityEffect, QMessageBox, QListView, QScrollArea, QWidgetAction, QApplication, QLabel, QGridLayout, \ QSlider, QGraphicsOpacityEffect, QMessageBox, QListView, QScrollArea, QWidgetAction, QApplication, QLabel, QGridLayout, \
QFileDialog, QListWidgetItem, QComboBox, QDialog, QAbstractItemView QFileDialog, QListWidgetItem, QComboBox, QDialog, QAbstractItemView, QSizePolicy
__dir__ = os.path.dirname(os.path.abspath(__file__)) __dir__ = os.path.dirname(os.path.abspath(__file__))
@ -227,6 +228,21 @@ class MainWindow(QMainWindow):
listLayout.addWidget(leftTopToolBoxContainer) listLayout.addWidget(leftTopToolBoxContainer)
# ================== Label List ================== # ================== Label List ==================
labelIndexListlBox = QHBoxLayout()
# Create and add a widget for showing current label item index
self.indexList = QListWidget()
self.indexList.setMaximumSize(40, 16777215) # limit max width
self.indexList.setEditTriggers(QAbstractItemView.NoEditTriggers) # no editable
self.indexList.itemSelectionChanged.connect(self.indexSelectionChanged)
self.indexList.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOff) # no scroll Bar
self.indexListDock = QDockWidget('No.', self)
self.indexListDock.setWidget(self.indexList)
self.indexListDock.setFeatures(QDockWidget.NoDockWidgetFeatures)
labelIndexListlBox.addWidget(self.indexListDock, 1)
# no margin between two boxes
labelIndexListlBox.setSpacing(0)
# Create and add a widget for showing current label items # Create and add a widget for showing current label items
self.labelList = EditInList() self.labelList = EditInList()
labelListContainer = QWidget() labelListContainer = QWidget()
@ -240,7 +256,7 @@ class MainWindow(QMainWindow):
self.labelListDock = QDockWidget(self.labelListDockName, self) self.labelListDock = QDockWidget(self.labelListDockName, self)
self.labelListDock.setWidget(self.labelList) self.labelListDock.setWidget(self.labelList)
self.labelListDock.setFeatures(QDockWidget.NoDockWidgetFeatures) self.labelListDock.setFeatures(QDockWidget.NoDockWidgetFeatures)
listLayout.addWidget(self.labelListDock) labelIndexListlBox.addWidget(self.labelListDock, 10) # label list is wider than index list
# enable labelList drag_drop to adjust bbox order # enable labelList drag_drop to adjust bbox order
# 设置选择模式为单选 # 设置选择模式为单选
@ -256,6 +272,17 @@ class MainWindow(QMainWindow):
# 触发放置 # 触发放置
self.labelList.model().rowsMoved.connect(self.drag_drop_happened) self.labelList.model().rowsMoved.connect(self.drag_drop_happened)
labelIndexListContainer = QWidget()
labelIndexListContainer.setLayout(labelIndexListlBox)
listLayout.addWidget(labelIndexListContainer)
# labelList indexList同步滚动
self.labelListBar = self.labelList.verticalScrollBar()
self.indexListBar = self.indexList.verticalScrollBar()
self.labelListBar.valueChanged.connect(self.move_scrollbar)
self.indexListBar.valueChanged.connect(self.move_scrollbar)
# ================== Detection Box ================== # ================== Detection Box ==================
self.BoxList = QListWidget() self.BoxList = QListWidget()
@ -766,6 +793,7 @@ class MainWindow(QMainWindow):
self.shapesToItemsbox.clear() self.shapesToItemsbox.clear()
self.labelList.clear() self.labelList.clear()
self.BoxList.clear() self.BoxList.clear()
self.indexList.clear()
self.filePath = None self.filePath = None
self.imageData = None self.imageData = None
self.labelFile = None self.labelFile = None
@ -1027,13 +1055,19 @@ class MainWindow(QMainWindow):
for shape in self.canvas.selectedShapes: for shape in self.canvas.selectedShapes:
shape.selected = False shape.selected = False
self.labelList.clearSelection() self.labelList.clearSelection()
self.indexList.clearSelection()
self.canvas.selectedShapes = selected_shapes self.canvas.selectedShapes = selected_shapes
for shape in self.canvas.selectedShapes: for shape in self.canvas.selectedShapes:
shape.selected = True shape.selected = True
self.shapesToItems[shape].setSelected(True) self.shapesToItems[shape].setSelected(True)
self.shapesToItemsbox[shape].setSelected(True) self.shapesToItemsbox[shape].setSelected(True)
index = self.labelList.indexFromItem(self.shapesToItems[shape]).row()
self.indexList.item(index).setSelected(True)
self.labelList.scrollToItem(self.currentItem()) # QAbstractItemView.EnsureVisible self.labelList.scrollToItem(self.currentItem()) # QAbstractItemView.EnsureVisible
# map current label item to index item and select it
index = self.labelList.indexFromItem(self.currentItem()).row()
self.indexList.scrollToItem(self.indexList.item(index))
self.BoxList.scrollToItem(self.currentBox()) self.BoxList.scrollToItem(self.currentBox())
if self.kie_mode: if self.kie_mode:
@ -1072,6 +1106,10 @@ class MainWindow(QMainWindow):
# item.setBackground(generateColorByText(shape.label)) # item.setBackground(generateColorByText(shape.label))
self.itemsToShapes[item] = shape self.itemsToShapes[item] = shape
self.shapesToItems[shape] = item self.shapesToItems[shape] = item
# add current label item index before label string
current_index = QListWidgetItem(str(self.labelList.count()))
current_index.setTextAlignment(Qt.AlignHCenter)
self.indexList.addItem(current_index)
self.labelList.addItem(item) self.labelList.addItem(item)
# print('item in add label is ',[(p.x(), p.y()) for p in shape.points], shape.label) # print('item in add label is ',[(p.x(), p.y()) for p in shape.points], shape.label)
@ -1105,6 +1143,7 @@ class MainWindow(QMainWindow):
del self.shapesToItemsbox[shape] del self.shapesToItemsbox[shape]
del self.itemsToShapesbox[item] del self.itemsToShapesbox[item]
self.updateComboBox() self.updateComboBox()
self.updateIndexList()
def loadLabels(self, shapes): def loadLabels(self, shapes):
s = [] s = []
@ -1156,6 +1195,13 @@ class MainWindow(QMainWindow):
# self.comboBox.update_items(uniqueTextList) # self.comboBox.update_items(uniqueTextList)
def updateIndexList(self):
self.indexList.clear()
for i in range(self.labelList.count()):
string = QListWidgetItem(str(i))
string.setTextAlignment(Qt.AlignHCenter)
self.indexList.addItem(string)
def saveLabels(self, annotationFilePath, mode='Auto'): def saveLabels(self, annotationFilePath, mode='Auto'):
# Mode is Auto means that labels will be loaded from self.result_dic totally, which is the output of ocr model # Mode is Auto means that labels will be loaded from self.result_dic totally, which is the output of ocr model
annotationFilePath = ustr(annotationFilePath) annotationFilePath = ustr(annotationFilePath)
@ -1211,6 +1257,10 @@ class MainWindow(QMainWindow):
# fix copy and delete # fix copy and delete
# self.shapeSelectionChanged(True) # self.shapeSelectionChanged(True)
def move_scrollbar(self, value):
self.labelListBar.setValue(value)
self.indexListBar.setValue(value)
def labelSelectionChanged(self): def labelSelectionChanged(self):
if self._noSelectionSlot: if self._noSelectionSlot:
return return
@ -1223,6 +1273,21 @@ class MainWindow(QMainWindow):
else: else:
self.canvas.deSelectShape() self.canvas.deSelectShape()
def indexSelectionChanged(self):
if self._noSelectionSlot:
return
if self.canvas.editing():
selected_shapes = []
for item in self.indexList.selectedItems():
# map index item to label item
index = self.indexList.indexFromItem(item).row()
item = self.labelList.item(index)
selected_shapes.append(self.itemsToShapes[item])
if selected_shapes:
self.canvas.selectShapes(selected_shapes)
else:
self.canvas.deSelectShape()
def boxSelectionChanged(self): def boxSelectionChanged(self):
if self._noSelectionSlot: if self._noSelectionSlot:
# self.BoxList.scrollToItem(self.currentBox(), QAbstractItemView.PositionAtCenter) # self.BoxList.scrollToItem(self.currentBox(), QAbstractItemView.PositionAtCenter)
@ -1517,6 +1582,7 @@ class MainWindow(QMainWindow):
if self.labelList.count(): if self.labelList.count():
self.labelList.setCurrentItem(self.labelList.item(self.labelList.count() - 1)) self.labelList.setCurrentItem(self.labelList.item(self.labelList.count() - 1))
self.labelList.item(self.labelList.count() - 1).setSelected(True) self.labelList.item(self.labelList.count() - 1).setSelected(True)
self.indexList.item(self.labelList.count() - 1).setSelected(True)
# show file list image count # show file list image count
select_indexes = self.fileListWidget.selectedIndexes() select_indexes = self.fileListWidget.selectedIndexes()
@ -2015,12 +2081,14 @@ class MainWindow(QMainWindow):
for shape in self.canvas.shapes: for shape in self.canvas.shapes:
shape.paintLabel = self.displayLabelOption.isChecked() shape.paintLabel = self.displayLabelOption.isChecked()
shape.paintIdx = self.displayIndexOption.isChecked() shape.paintIdx = self.displayIndexOption.isChecked()
self.canvas.repaint()
def togglePaintIndexOption(self): def togglePaintIndexOption(self):
self.displayLabelOption.setChecked(False) self.displayLabelOption.setChecked(False)
for shape in self.canvas.shapes: for shape in self.canvas.shapes:
shape.paintLabel = self.displayLabelOption.isChecked() shape.paintLabel = self.displayLabelOption.isChecked()
shape.paintIdx = self.displayIndexOption.isChecked() shape.paintIdx = self.displayIndexOption.isChecked()
self.canvas.repaint()
def toogleDrawSquare(self): def toogleDrawSquare(self):
self.canvas.setDrawingShapeToSquare(self.drawSquaresOption.isChecked()) self.canvas.setDrawingShapeToSquare(self.drawSquaresOption.isChecked())
@ -2254,6 +2322,7 @@ class MainWindow(QMainWindow):
self.itemsToShapesbox.clear() # ADD self.itemsToShapesbox.clear() # ADD
self.shapesToItemsbox.clear() self.shapesToItemsbox.clear()
self.labelList.clear() self.labelList.clear()
self.indexList.clear()
self.BoxList.clear() self.BoxList.clear()
self.result_dic = [] self.result_dic = []
self.result_dic_locked = [] self.result_dic_locked = []
@ -2665,6 +2734,7 @@ class MainWindow(QMainWindow):
def undoShapeEdit(self): def undoShapeEdit(self):
self.canvas.restoreShape() self.canvas.restoreShape()
self.labelList.clear() self.labelList.clear()
self.indexList.clear()
self.BoxList.clear() self.BoxList.clear()
self.loadShapes(self.canvas.shapes) self.loadShapes(self.canvas.shapes)
self.actions.undo.setEnabled(self.canvas.isShapeRestorable) self.actions.undo.setEnabled(self.canvas.isShapeRestorable)
@ -2674,6 +2744,7 @@ class MainWindow(QMainWindow):
for shape in shapes: for shape in shapes:
self.addLabel(shape) self.addLabel(shape)
self.labelList.clearSelection() self.labelList.clearSelection()
self.indexList.clearSelection()
self._noSelectionSlot = False self._noSelectionSlot = False
self.canvas.loadShapes(shapes, replace=replace) self.canvas.loadShapes(shapes, replace=replace)
print("loadShapes") # 1 print("loadShapes") # 1