class PropertyEditor(QWidget):
# Signals
insertionModeStarted = pyqtSignal(str)
insertionModeEnded = pyqtSignal()
insertionPropertiesChanged = pyqtSignal(object)
editPropertiesChanged = pyqtSignal(object)
def __init__(self, config, parent=None):
QWidget.__init__(self, parent)
self._class_config = {}
self._class_items = {}
self._class_prototypes = {}
self._attribute_handlers = {}
self._handler_factory = AttributeHandlerFactory()
self._setupGUI()
self._parea.setGeometry(0, 0, 200, 0)
# (快捷键,button)
self.shortcut2button = {}
# (button,快捷键)
self.label2shortcut = {}
# Add label classes from config
for label in config:
self.addLabelClass(label)
self.image_path = None
def addLabelClassByPath(self, configs_path):
# 读配置文件
with open(configs_path, 'r') as f:
configs = json5.load(f)
# 写入当前配置文件的路径
direct = os.path.dirname(sys.argv[0])
with open(os.path.join(direct, 'sloth.txt'), 'w') as f:
f.write(configs_path)
self._parea.setGeometry(0, 0, 200, 0)
for temp_json in configs:
self.addLabelClass(temp_json)
# 注册
self._register('inserter', temp_json['attributes']['class'],
temp_json['inserter'])
self._register('item', temp_json['attributes']['class'],
temp_json['item'])
# add_txt的下拉框里也要添加
self.combo_box.addItem(temp_json['attributes']['class'])
self.items.append(temp_json['attributes']['class'])
cf.LABELS.append(temp_json)
def onModelChanged(self, new_model):
attrs = set([
k for k, v in self._attribute_handlers.items()
if v.autoAddEnabled()
])
if len(attrs) > 0:
start = time.time()
attr2vals = {}
for item in new_model.iterator(AnnotationModelItem):
for attr in attrs:
if attr in item:
if attr not in attr2vals:
attr2vals[attr] = set((item[attr], ))
else:
attr2vals[attr] |= set((item[attr], ))
diff = time.time() - start
LOG.info("Extracted annotation values from model in %.2fs" % diff)
for attr, vals in attr2vals.items():
h = self._attribute_handlers[attr]
for val in vals:
h.addValue(val, True)
# 设置右键菜单所在位置
def showContextMenu(self, label_class):
self.label_menu[label_class].exec_(QCursor.pos())
# 删除所有的item
def remove_all_item(self):
for v in self.shortcut2button.values():
v.setShortcut(QKeySequence())
self.shortcut2button.clear()
self.label2shortcut.clear()
self.label_menu.clear()
self.label_action.clear()
self._class_config.clear()
self.combo_box.clear()
self.items.clear()
temp_dict = copy.copy(self._class_buttons)
for k, v in temp_dict.items():
self._classbox_layout.removeWidget(v)
# 下面这句很重要，不然相当于没删
self._class_buttons[k].deleteLater()
self._class_buttons.clear()
cf.LABELS.clear()
self._parea.setGeometry(0, 0, 200, 60)
# 删除标签
def remove_item(self, label_class):
"""
删除标签
:param label_class: 删除的标签名字
"""
try:
# 删除快捷键
self.endInsertionMode()
shortcut = self.label2shortcut[label_class]
if shortcut in self.shortcut2button and \
self.shortcut2button[shortcut] is not None and \
self.shortcut2button[shortcut] == self._class_buttons[label_class]:
self.shortcut2button[shortcut].setShortcut(QKeySequence())
del self.shortcut2button[shortcut]
# 删除菜单
del self.label_menu[label_class]
# 删除菜单的动作
del self.label_action[label_class]
# 删除视图中的按钮
self._classbox_layout.removeWidget(
self._class_buttons[label_class])
self._class_buttons[label_class].deleteLater()
self._class_buttons[label_class] = None
del self._class_config[label_class]
# 写回json
direct = os.path.dirname(sys.argv[0])
with open(os.path.join(direct, 'sloth.txt'), 'r') as f:
label_path = f.read()
try:
with open(label_path, 'r') as f:
temp = json5.load(f)
for i, current_json in enumerate(temp):
if current_json['attributes']['class'] == label_class:
temp.remove(current_json)
# 遍历combo box 找到要删的
for i in range(len(self.combo_box)):
current_label = self.combo_box.itemText(i)
if current_label == label_class:
print('removed', label_class)
self.combo_box.removeItem(i)
self.items.pop(i)
break
self._class_buttons.pop(label_class)
break
with open(label_path, 'w') as f:
json5.dump(temp,
f,
quote_keys=True,
trailing_commas=False,
indent=4,
separators=(',', ': '),
sort_keys=True,
ensure_ascii=False)
self._parea.setGeometry(
0, 0, 200, max(self._parea.geometry().height() - 40, 60))
except Exception as e:
print(e)
except Exception as e:
print(e)
# 添加标签
def addLabelClass(self, label_config):
"""
添加标签
:param label_config: 标签的json
"""
# Check label configuration
if 'attributes' not in label_config:
raise ImproperlyConfigured("Label with no 'attributes' dict found")
attrs = label_config['attributes']
if 'class' not in attrs:
raise ImproperlyConfigured("Labels must have an attribute 'class'")
label_class = attrs['class']
if label_class in self._class_config:
raise ImproperlyConfigured(
"Label with class '%s' defined more than once" % label_class)
# Store config
self._class_config[label_class] = label_config
# Parse configuration and create handlers and item
self.parseConfiguration(label_class, label_config)
# Add label class button
button_text = label_config['text']
button = QPushButton(button_text)
button.setCheckable(True)
button.setFlat(True)
button.clicked.connect(bind(self.onClassButtonPressed, label_class))
self._class_buttons[label_class] = button
self._parea.setGeometry(0, 0, 200,
self._parea.geometry().height() + 40)
self._classbox_layout.addWidget(button)
# 添加右键菜单
self.label_menu[label_class] = QtGui.QMenu(self)
self.label_action[label_class] = self.label_menu[
label_class].addAction('删除')
self.label_action[label_class].triggered.connect(
bind(self.remove_item, label_class))
self._class_buttons[label_class].setContextMenuPolicy(
QtCore.Qt.CustomContextMenu)
self._class_buttons[label_class].customContextMenuRequested.connect(
bind(self.showContextMenu, label_class))
# Add hotkey
if 'hotkey' in label_config:
# 快捷键
hotkey = label_config['hotkey']
# 快捷键已经存在，那就去掉原来的
if hotkey in self.shortcut2button and self.shortcut2button[
hotkey] is not None:
self.shortcut2button[hotkey].setShortcut(QKeySequence())
# 设置快捷键
button.setShortcut(QKeySequence(hotkey))
self.shortcut2button[hotkey] = button
self.label2shortcut[label_class] = hotkey
def parseConfiguration(self, label_class, label_config):
attrs = label_config['attributes']
# Add prototype item for insertion
self._class_items[label_class] = AnnotationModelItem(
{'class': label_class})
# Create attribute handler widgets or update their values
for attr, vals in attrs.items():
if attr in self._attribute_handlers:
self._attribute_handlers[attr].updateValues(vals)
else:
handler = self._handler_factory.create(attr, vals)
if handler is None:
self._class_items[label_class][attr] = vals
else:
self._attribute_handlers[attr] = handler
for attr in attrs:
if attr in self._attribute_handlers:
self._class_items[label_class].update(
self._attribute_handlers[attr].defaults())
def getHandler(self, attribute):
if attribute in self._attribute_handlers:
return self._attribute_handlers[attribute]
else:
return None
def getLabelClassAttributes(self, label_class):
return self._class_config[label_class]['attributes'].keys()
def onClassButtonPressed(self, label_class):
if self._class_buttons[label_class].isChecked():
self.startInsertionMode(label_class)
else:
self.endInsertionMode()
def startInsertionMode(self, label_class):
self.endInsertionMode(False)
for lc, button in self._class_buttons.items():
button.setChecked(lc == label_class)
LOG.debug("Starting insertion mode for %s" % label_class)
self._label_editor = LabelEditor([self._class_items[label_class]],
self, True)
# self._layout.insertWidget(1, self._label_editor, 0)
self.insertionModeStarted.emit(label_class)
def endInsertionMode(self, uncheck_buttons=True):
if self._label_editor is not None:
LOG.debug("Ending insertion/edit mode")
self._label_editor.hide()
# self._layout.removeWidget(self._label_editor)
self._label_editor = None
if uncheck_buttons:
self.uncheckAllButtons()
self.insertionModeEnded.emit()
def uncheckAllButtons(self):
for lc, button in self._class_buttons.items():
button.setChecked(False)
def markEditButtons(self, label_classes):
for lc, button in self._class_buttons.items():
button.setFlat(lc not in label_classes)
def currentEditorProperties(self):
if self._label_editor is None:
return None
else:
return self._label_editor.currentProperties()
def startEditMode(self, model_items):
# If we're in insertion mode, ignore empty edit requests
if self._label_editor is not None and self._label_editor.insertionMode() \
and len(model_items) == 0:
return
self.endInsertionMode()
LOG.debug("Starting edit mode for items: %s" % model_items)
self._label_editor = LabelEditor(model_items, self)
self.markEditButtons(self._label_editor.labelClasses())
self._layout.insertWidget(1, self._label_editor, 0)
# 添加txt
def add_txt(self):
if not Main.isConfig(Main.get_json()):
QMessageBox.warning(self, "Warning", '当前的配置文件错误或者为空,无法添加txt',
QMessageBox.Ok)
defect = self.combo_box.currentText()
if defect is None or defect == '':
return
dir_path = QFileDialog.getExistingDirectory(self)
Main.write_txt(dir_path, {defect}, 'defect')
# 选择图片
def select_image(self):
image_types = [
'*.jpg', '*.bmp', '*.png', '*.pgm', '*.ppm', '*.tiff', '*.tif',
'*.gif'
]
format_str = ' '.join(image_types)
fname = QFileDialog.getOpenFileName(
self, "select training source", '.',
"Media files (%s)" % (format_str, ))
if fname is None or fname == '':
return
self.image_path = os.path.abspath(fname)
self._image_label.setText(os.path.basename(self.image_path))
# 获得图片路径对应的json路径
def image2json(self, path):
temp = path.split('.')
return ''.join(temp[:-1]) + '.json'
# 获得图片路径转成的训练图片路径
def image2cpimage(self, path, id, length):
length = max(length, 5)
temp = path.split('.')
return ''.join(temp[:-1]) + str(id).zfill(length) + '.' + temp[-1]
# 判断是否包含瑕疵
def contains_defect(self, annotations, defect_type):
defects = set()
for annotation in annotations:
if 'class' in annotation:
defects.add(annotation['class'])
return defect_type.issubset(defects)
# 生成训练数据
def generate(self):
if not Main.isConfig(Main.get_json()):
QMessageBox.warning(self, "Warning", '配置文件错误或者为空，无法生成训练数据',
QMessageBox.Ok)
return
a = TrainDialog(self)
a.exec_()
# 从labeltool中设置搜索按钮
def setFunction(self, func):
self._search_btn.clicked.connect(func)
# 获得关键字
def get_key_word(self):
key_word = self._key_word.text()
if key_word is None or key_word == '':
key_word = self._key_word.placeholderText()
return key_word
# 获得文件类型
def get_extension(self):
extension = self._extension.text()
# 为空则用默认的，否则用输入的
if extension is None or extension == '':
extension = self._extension.placeholderText()
return extension
# 返回一个含有权限类型的list
def get_attributes_type(self):
'''
'Rect':('sloth.items.RectItem','sloth.items.RectItemInserter'),
'Point':('sloth.items.PointItem','sloth.items.PointItemInserter'),
'Polygon':('sloth.items.PolygonItem','sloth.items.PolygonItemInserter')
'''
return ['Rect', 'Point', 'Polygon']
# 写回json
def rewrite_json(self, temp_json):
# json所在的txt
direct = os.path.dirname(sys.argv[0])
with open(os.path.join(direct, 'sloth.txt'), 'r') as f:
label_path = f.read()
try:
# 读取旧json
with open(label_path, 'r') as f:
temp = json5.load(f)
# 追加我们要写入的json
temp.append(temp_json)
# 写入
with open(label_path, 'w') as f:
json5.dump(temp,
f,
quote_keys=True,
trailing_commas=False,
indent=4,
separators=(',', ': '),
sort_keys=True,
ensure_ascii=False)
except Exception as e:
print(e)
# 添加标签
def add_attributes(self):
if not Main.isConfig(Main.get_json()):
QMessageBox.warning(self, "Warning", '当前配置文件错误或者为空，不能添加标签',
QMessageBox.Ok)
return
# 转换dict
type_dict = {
'Rect': ('sloth.items.RectItem', 'sloth.items.RectItemInserter'),
'Point':
('sloth.items.PointItem', 'sloth.items.PointItemInserter'),
'Polygon':
('sloth.items.PolygonItem', 'sloth.items.PolygonItemInserter')
}
# 获取添加的标签信息
attributes = {'class': self.attributes_LineEdit.text()}
attributes_item, attributes_inserter = type_dict[
self.attributes_type.currentText()]
attributes_hotkey = self.hotkey.text()
attributes_text = self.text_LineEdit.text()
global brush2idx
brush_idx = str(brush2idx[self.brush_combo_box.currentText()])
temp_json = {
'attributes': attributes,
'inserter': attributes_inserter,
'item': attributes_item,
'color': ','.join(map(str, self.color_info)),
'brush': brush_idx,
'text': attributes_text
}
# 快捷键
if attributes_hotkey is not None and attributes_hotkey != '':
temp_json['hotkey'] = attributes_hotkey
print(temp_json)
try:
# 加入标签
self.addLabelClass(temp_json)
print(self._class_buttons.keys())
# 注册
self._register('inserter', temp_json['attributes']['class'],
temp_json['inserter'])
self._register('item', temp_json['attributes']['class'],
temp_json['item'])
# add_txt的下拉框里也要添加
self.combo_box.addItem(temp_json['attributes']['class'])
self.items.append(temp_json['attributes']['class'])
cf.LABELS.append(temp_json)
# 写回json
self.rewrite_json(temp_json)
except Exception as e:
print(e)
# 颜色对话框
def color_dialog(self):
col = QtGui.QColorDialog.getColor()
if col.isValid():
self.color_label.setStyleSheet("QWidget { background-color: %s }" %
col.name())
self.color_info = col.getRgb()[:-1]
# 设置控件的隐藏状态
def component_visible(self, component_name, state):
if component_name == '添加标签':
self._group_box_add_label.setVisible(state)
elif component_name == 'add_txt':
self._group_box_add_txt.setVisible(state)
elif component_name == 'add_files':
self._group_box_add_files.setVisible(state)
def _setupGUI(self):
self._class_buttons = {}
self.label_menu = {}
self.label_action = {}
self._label_editor = None
# Label class buttons
self._parea = QGroupBox("Labels")
self._classbox = QScrollArea()
self._classbox_layout = FloatingLayout()
self._parea.setLayout(self._classbox_layout)
self._parea.setGeometry(0, 0, 200, 200)
self._classbox.setWidget(self._parea)
self._classbox.setGeometry(0, 0, 100, 100)
# 添加txt模块
self.combo_box = QComboBox()
self._group_box_add_txt = QGroupBox('add_txt', self)
self._group_box_add_txt_layout = QVBoxLayout()
self._group_box_add_txt.setLayout(self._group_box_add_txt_layout)
temp = cf.LABELS
self.items = []
# 获取所有的标签
for i in temp:
self.items.append(i['attributes']['class'])
# 假如下拉框
self.combo_box.addItems(self.items)
self.add_txt_btn = QPushButton('add txt')
self.add_txt_btn.clicked.connect(self.add_txt)
# 加入下拉框和按钮
self._group_box_add_txt_layout.addWidget(self.combo_box, 0)
self._group_box_add_txt_layout.addWidget(self.add_txt_btn, 1)
# 根据关键字搜索图片模块
self._group_box_add_files = QGroupBox('add files', self)
# 文件名包含的
self._key_word = QLineEdit('')
self._key_word.setPlaceholderText('Second')
# 文件类型
self._extension = QLineEdit('')
self._extension.setPlaceholderText('bmp')
self._search_btn = QPushButton('search files')
self._group_box_add_files_layout = QVBoxLayout()
# 加入控件
self._group_box_add_files_layout.addWidget(self._key_word, 0)
self._group_box_add_files_layout.addWidget(self._extension, 1)
self._group_box_add_files_layout.addWidget(self._search_btn, 2)
self._group_box_add_files.setLayout(self._group_box_add_files_layout)
# 添加标签模块
self._group_box_add_label = QGroupBox("添加标签", self)
self._add_label_group_layout = QVBoxLayout()
self._group_box_add_label.setLayout(self._add_label_group_layout)
# 标签的class
self.attributes_LineEdit = QLineEdit('')
self.attributes_LineEdit.setPlaceholderText('attributes')
# 标签画出来的类型
self.attributes_type = QComboBox()
self.attributes_type.addItems(self.get_attributes_type())
# 快捷键，目前设置了只允许一个键
self.hotkey = QLineEdit('')
self.hotkey.setPlaceholderText('hotkey')
self.regx = QRegExp("[a-z0-9]$")
self.validator = QRegExpValidator(self.regx, self.hotkey)
self.hotkey.setValidator(self.validator)
# 标签显示
self.text_LineEdit = QLineEdit('')
self.text_LineEdit.setPlaceholderText('text')
# 颜色
color = QtGui.QColor(0, 0, 0)
self.color_label = QtGui.QWidget()
self.color_label.setStyleSheet("QWidget { background-color: %s }" %
color.name())
self.color_info = [0, 0, 0]
self.color_layout = QtGui.QHBoxLayout()
self.color_btn = QPushButton('选择颜色')
self.color_btn.clicked.connect(self.color_dialog)
self.color_layout.addWidget(self.color_label)
self.color_layout.addWidget(self.color_btn)
# 笔刷
global brush2idx
self.brush_combo_box = QComboBox()
self.brush_combo_box.addItems(list(brush2idx.keys()))
# 按钮
self.attributes_add_btn = QPushButton('添加标签')
self.attributes_add_btn.clicked.connect(self.add_attributes)
# 加入控件
self._add_label_group_layout.addWidget(self.attributes_LineEdit, 0)
self._add_label_group_layout.addWidget(self.attributes_type, 1)
self._add_label_group_layout.addWidget(self.hotkey, 2)
self._add_label_group_layout.addWidget(self.text_LineEdit, 3)
self._label_widget = QWidget()
self._label_widget.setLayout(self.color_layout)
self._add_label_group_layout.addWidget(self._label_widget, 4)
self._add_label_group_layout.addWidget(self.brush_combo_box, 5)
self._add_label_group_layout.addWidget(self.attributes_add_btn, 6)
# 生成训练数据按钮
self._file_button = QPushButton('生成训练数据')
self._file_button.clicked.connect(self.generate)
# Global widget
self._layout = MyVBoxLayout()
self.setLayout(self._layout)
self._layout.addWidget(self._classbox, 1)
self._layout.insertWidget(-1, self._group_box_add_label, 1)
self._layout.insertWidget(-1, self._group_box_add_txt, 1)
self._layout.insertWidget(-1, self._group_box_add_files, 1)
self._layout.insertWidget(-1, self._file_button, 1)
