def new_shape(self): """Pop-up and give focus to the label editor. position MUST be in global coordinates. """ BB = QDialogButtonBox x1, y1, x2, y2 = int(self.canvas.line[0].x()), int( self.canvas.line[0].y()), int(self.canvas.line[1].x()), int( self.canvas.line[1].y()) image = self.image[y1:y2, x1:x2] self.labelDialog = LabelDialog(parent=self, image=image) self.labelDialog.show()
def __init__(self, camera=None, path=None): super(DatabaseEditor, self).__init__() self.setupUi(self) self.camera = camera self.stream_enabled = False self.dirty = False self.frame = None self.items_to_shapes = {} self.shapes_to_items = {} self.prev_label_text = '' self.last_label = None self.shapes = [] self.path = path self._no_selection_slot = False self.current_component = '' self.current_record = '' self.database_handler = DatabaseHandler(path) self.label_list = self.database_handler.classes.copy() self.label_dialog = LabelDialog(parent=self, listItem=self.label_list) self.scrollBars = { Qt.Vertical: self.scrollArea.verticalScrollBar(), Qt.Horizontal: self.scrollArea.horizontalScrollBar() } # context menus delete_record_action = Action(self, 'Delete record', self.delete_record, enabled=True) add_record_action = Action(self, 'Add record', self.add_record, enabled=True) self.record_menu = QMenu() add_actions(self.record_menu, (delete_record_action, add_record_action)) delete_component_action = Action(self, 'Delete component', self.delete_component, enabled=True) add_component_action = Action(self, 'Add component', self.add_component, enabled=True) add_component_image_action = Action(self, 'Add components image', self.add_component_image, enabled=True) self.component_menu = QMenu() add_actions(self.component_menu, (delete_component_action, add_component_action, add_component_image_action)) delete_shape_action = Action(self, 'Delete rectangle', self.delete_shape, enabled=True) self.rectangle_menu = QMenu() add_actions(self.rectangle_menu, delete_shape_action) self.modeEdit.setChecked(True) self.modeSelect.setChecked(False) self.labelCoordinates = QLabel('') self.statusBar().addPermanentWidget(self.labelCoordinates) self.connect() self.display_classes()
def save_labels(self): component = self.current_component record = self.current_record self.label_dialog = LabelDialog(parent=self, listItem=self.database_handler.classes) if record: self.database_handler.edit_record(component, record, self.frame, self.shapes) else: if component: text = component else: text = self.label_dialog.popUp(self.last_label) self.database_handler.add_record(text, self.frame, self.shapes) self.display_classes() self.set_clean()
def new_shape(self): if len(self.label_list): self.label_dialog = LabelDialog(parent=self, listItem=self.label_list) text = self.label_dialog.popUp(text=self.prev_label_text) self.last_label = text if text is not None: self.prev_label_text = text color = generate_color_by_text(text) shape = self.canvas.setLastLabel(text, color, color) self.add_label(shape) self.shapes.append(shape) if text not in self.label_list: self.label_list.append(text) self.set_dirty() else: self.canvas.resetAllLines() self.mode_edit(False)
class CaptureDialog(QDialog): def __init__(self, parent=None): QDialog.__init__(self, parent) self.parent = parent self.data = None self.init_ui() def init_ui(self): self.centralWidget = QtWidgets.QFrame(self) self.centralWidget.setObjectName("centralWidget") self.centralWidget.setGeometry(QtCore.QRect(0, 0, 600, 600)) self.centralWidget.setFrameShape(QtWidgets.QFrame.Box) self.frame_2 = QtWidgets.QFrame(self.centralWidget) self.frame_2.setObjectName("frame_2") self.frame_2.setGeometry(QtCore.QRect(10, 10, 580, 580)) self.frame_2.setFrameShape(QtWidgets.QFrame.Box) self.canvas = Canvas(self.frame_2) self.canvas.setGeometry(QtCore.QRect(0, 0, 580, 580)) self.canvas.setEnabled(True) self.canvas.setFocus(True) self.canvas.setDrawingShapeToSquare(False) self.canvas.restoreCursor() self.canvas.mode = self.canvas.CREATE self.image = None self.canvas.show() self.canvas.newShape.connect(self.new_shape) def new_shape(self): """Pop-up and give focus to the label editor. position MUST be in global coordinates. """ BB = QDialogButtonBox x1, y1, x2, y2 = int(self.canvas.line[0].x()), int( self.canvas.line[0].y()), int(self.canvas.line[1].x()), int( self.canvas.line[1].y()) image = self.image[y1:y2, x1:x2] self.labelDialog = LabelDialog(parent=self, image=image) self.labelDialog.show()
class DatabaseEditor(QMainWindow, designer.Ui_MainWindow): close_event = pyqtSignal() def __init__(self, camera=None, path=None): super(DatabaseEditor, self).__init__() self.setupUi(self) self.camera = camera self.stream_enabled = False self.dirty = False self.frame = None self.items_to_shapes = {} self.shapes_to_items = {} self.prev_label_text = '' self.last_label = None self.shapes = [] self.path = path self._no_selection_slot = False self.current_component = '' self.current_record = '' self.database_handler = DatabaseHandler(path) self.label_list = self.database_handler.classes.copy() self.label_dialog = LabelDialog(parent=self, listItem=self.label_list) self.scrollBars = { Qt.Vertical: self.scrollArea.verticalScrollBar(), Qt.Horizontal: self.scrollArea.horizontalScrollBar() } # context menus delete_record_action = Action(self, 'Delete record', self.delete_record, enabled=True) add_record_action = Action(self, 'Add record', self.add_record, enabled=True) self.record_menu = QMenu() add_actions(self.record_menu, (delete_record_action, add_record_action)) delete_component_action = Action(self, 'Delete component', self.delete_component, enabled=True) add_component_action = Action(self, 'Add component', self.add_component, enabled=True) add_component_image_action = Action(self, 'Add components image', self.add_component_image, enabled=True) self.component_menu = QMenu() add_actions(self.component_menu, (delete_component_action, add_component_action, add_component_image_action)) delete_shape_action = Action(self, 'Delete rectangle', self.delete_shape, enabled=True) self.rectangle_menu = QMenu() add_actions(self.rectangle_menu, delete_shape_action) self.modeEdit.setChecked(True) self.modeSelect.setChecked(False) self.labelCoordinates = QLabel('') self.statusBar().addPermanentWidget(self.labelCoordinates) self.connect() self.display_classes() def delete_shape(self): self.remove_label(self.canvas.deleteSelected()) self.save_labels() def remove_label(self, shape): if shape is None: return self.shapes.remove(shape) item = self.shapes_to_items[shape] self.rectangleList.takeItem(self.rectangleList.row(item)) del self.shapes_to_items[shape] del self.items_to_shapes[item] def add_component_image(self): image = self.frame component = self.get_selected_component().text() self.database_handler.add_ideal_image(image, component) msg = 'Class\'s image added! You could see it at the main window of the program.' QMessageBox.information(self, 'Information', msg, QMessageBox.Ok) def delete_component(self): component = self.get_selected_component() if not component: return component = component.text() self.database_handler.delete_class(component) self.clear() self.database_handler.index_ideal_images() self.display_classes() def add_component(self): self.clear() msg = 'Creating new class. Please add at least 15 images and mark them out!' QMessageBox.information(self, 'Information', msg, QMessageBox.Ok) component = self.label_dialog.popUp() if not component: return self.database_handler.add_class(component) self.display_classes() def new_shape(self): if len(self.label_list): self.label_dialog = LabelDialog(parent=self, listItem=self.label_list) text = self.label_dialog.popUp(text=self.prev_label_text) self.last_label = text if text is not None: self.prev_label_text = text color = generate_color_by_text(text) shape = self.canvas.setLastLabel(text, color, color) self.add_label(shape) self.shapes.append(shape) if text not in self.label_list: self.label_list.append(text) self.set_dirty() else: self.canvas.resetAllLines() self.mode_edit(False) def display_classes(self): self.componentList.clear() self.label_list = self.database_handler.classes.copy() self.componentList.addItems(self.label_list) def display_records(self): if not self.may_continue(): return self.set_clean() self.clear() item = self.componentList.selectedItems()[0] if not item: return self.recordList.clear() self.current_component = item.text() if item.text() in self.database_handler.records.keys(): for record in self.database_handler.records[item.text()]: self.recordList.addItem("{} №{}".format( record.date, record.number)) def load_record(self): if not self.may_continue(): return self.set_clean() self.clear_labels() item = self.recordList.selectedItems()[0] if not item: return self.current_record = item.text() if self.stream_enabled: self._stop_video() component = self.componentList.selectedItems()[0].text() filename = item.text().replace('№', '') f = filename.split(' ') filename = '{}{:04d}'.format(f[0], int(f[1])) path = os.path.join(self.path, component, 'records', filename) img_path = path + '.jpg' txt_path = path + '.txt' self.frame = cv2.imread(img_path) if self.frame is None: return self.frame = cv2.cvtColor(self.frame, cv2.COLOR_BGR2RGB) self.canvas.loadPixmap( QPixmap.fromImage(qimage2ndarray.array2qimage(self.frame))) self.canvas.adjustSize() with open(txt_path, 'r') as bndboxes_file: for box in bndboxes_file: index, xcen, ycen, w, h = box.strip().split(' ') label = self.database_handler.classes[int(index)] xmin, ymin, xmax, ymax = yolo2points(xcen, ycen, w, h, self.frame.shape[1], self.frame.shape[0]) points = [(xmin, ymin), (xmax, ymin), (xmax, ymax), (xmin, ymax)] shape = Shape(label=label) for x, y in points: x, y, snapped = self.canvas.snapPointToCanvas(x, y) if snapped: self.set_dirty() shape.addPoint(QPointF(x, y)) shape.difficult = False shape.fill_color = generate_color_by_text(label) shape.line_color = generate_color_by_text(label) shape.close() self.shapes.append(shape) self.add_label(shape) self.canvas.loadShapes(self.shapes) def delete_record(self): record = self.get_selected_record() component = self.get_selected_component() if not record or not component: return component = component.text() filename = record.text().replace('№', '') f = filename.split(' ') filename = '{}{:04d}'.format(f[0], int(f[1])) path = os.path.join(self.path, component, 'records', filename) os.remove(path + '.jpg') os.remove(path + '.txt') for i, o in enumerate(self.database_handler.records[component]): if o.image == filename + '.txt': del self.database_handler.records[component][i] break self.clear() self.display_records() def add_record(self): component = self.get_selected_component() if not component: return component = component.text() if not self.stream_enabled: self.database_handler.add_record(component, self.frame, self.shapes) self.display_records() def clear_labels(self): self.last_label = None self.items_to_shapes.clear() self.shapes_to_items.clear() self.prev_label_text = '' self.shapes.clear() if self.rectangleList.count(): self.rectangleList.clear() def add_label(self, shape): item = HashableQListWidgetItem(shape.label) item.setFlags(item.flags() | Qt.ItemIsUserCheckable) item.setCheckState(Qt.Checked) item.setBackground(generate_color_by_text(shape.label)) self.items_to_shapes[item] = shape self.shapes_to_items[shape] = item self.rectangleList.addItem(item) def edit_label(self): item = self.get_current_rectangle() if not item: return text = self.label_dialog.popUp(item.text()) if text is not None: item.setText(text) item.setBackground(generate_color_by_text(text)) self.set_dirty() def save_labels(self): component = self.current_component record = self.current_record self.label_dialog = LabelDialog(parent=self, listItem=self.database_handler.classes) if record: self.database_handler.edit_record(component, record, self.frame, self.shapes) else: if component: text = component else: text = self.label_dialog.popUp(self.last_label) self.database_handler.add_record(text, self.frame, self.shapes) self.display_classes() self.set_clean() def connect(self): self.camera.new_frame.connect(self._on_new_frame) self.canvas.shapeMoved.connect(self.set_dirty) self.canvas.selectionChanged.connect(self.canvas_shape_selected) self.canvas.newShape.connect(self.new_shape) self.shotButton.clicked.connect(self._stop_video) self.saveButton.clicked.connect(self.save_labels) self.modeSelect.triggered.connect(partial(self.mode_edit, False)) self.modeEdit.triggered.connect(partial(self.mode_edit, True)) self.addClass.triggered.connect(self.add_component) self.addImage.triggered.connect(self.add_component_image) self.componentList.customContextMenuRequested.connect( self.class_menu_popup) self.componentList.itemClicked.connect(self.display_records) # self.componentList.itemSelectionChanged.connect(self.clear) self.recordList.itemClicked.connect(self.load_record) # self.recordList.itemSelectionChanged.connect(self.clear_shapes) self.recordList.customContextMenuRequested.connect( self.record_menu_popup) self.rectangleList.itemDoubleClicked.connect(self.edit_label) self.rectangleList.itemActivated.connect(self.select_shape) self.rectangleList.itemSelectionChanged.connect(self.select_shape) self.rectangleList.itemChanged.connect(self.change_shape) self.rectangleList.customContextMenuRequested.connect( self.rectangle_menu_popup) def change_shape(self, item): shape = self.items_to_shapes[item] label = item.text() if label != shape.label: shape.label = label shape.line_color = generate_color_by_text(label) self.set_dirty() else: self.canvas.setShapeVisible(shape, True) def select_shape(self): item = self.get_current_rectangle() if item and self.canvas.editing(): self._no_selection_slot = True self.canvas.selectShape(self.items_to_shapes[item]) # helpers def get_current_rectangle(self): items = self.rectangleList.selectedItems() if items: return items[0] return None def get_selected_record(self): records = self.recordList.selectedItems() if records: return records[0] return None def get_selected_component(self): component = self.componentList.selectedItems() if component: return component[0] return None def set_clean(self): self.dirty = False def set_dirty(self): self.dirty = True def clear(self): self.last_label = None if not self.stream_enabled: self._stop_video() self.mode_edit(True) self.shapes.clear() self.shapes_to_items.clear() self.items_to_shapes.clear() if self.recordList.count(): self.recordList.clear() if self.rectangleList.count(): self.rectangleList.clear() self.canvas.resetAllLines() self.canvas.adjustSize() def clear_shapes(self): self.shapes.clear() self.shapes_to_items.clear() self.items_to_shapes.clear() if self.rectangleList.count(): self.rectangleList.clear() self.canvas.resetAllLines() self.canvas.adjustSize() def discard_changes_dialog(self): yes, no, cancel = QMessageBox.Yes, QMessageBox.No, QMessageBox.Cancel message = 'You have unsaved changes, would you like to save them and proceed?\nClick "No" to undo all changes.' return QMessageBox.warning(self, 'Attention', message, cancel | no | yes) def may_continue(self): if not self.dirty: return True else: action = self.discard_changes_dialog() if action == QMessageBox.No: return True elif action == QMessageBox.Yes: self.save_labels() return True else: return False # signal functions def canvas_shape_selected(self): if self._no_selection_slot: self._no_selection_slot = False else: shape = self.canvas.selectedShape if shape: self.shapes_to_items[shape].setSelected(True) else: self.rectangleList.clearSelection() def record_menu_popup(self, point): self.record_menu.exec_(self.recordList.mapToGlobal(point)) def class_menu_popup(self, point): self.component_menu.exec_(self.componentList.mapToGlobal(point)) def rectangle_menu_popup(self, point): self.rectangle_menu.exec_(self.rectangleList.mapToGlobal(point)) def mode_edit(self, edit=True): self.modeEdit.setChecked(edit) self.modeSelect.setChecked(not edit) self.canvas.setEditing(edit) def _stop_video(self): if self.stream_enabled: self.stream_enabled = False self.shotButton.setText( QCoreApplication.translate("MainWindow", "Включить видео")) else: self.stream_enabled = True self.shotButton.setText( QCoreApplication.translate("MainWindow", "Сделать снимок")) if len(self.shapes): self.clear() def _on_new_frame(self, frame): if self.stream_enabled: self.frame = cv2.cvtColor(frame.copy(), cv2.COLOR_BGR2RGB) self.canvas.loadPixmap( QPixmap.fromImage(qimage2ndarray.array2qimage(self.frame))) def closeEvent(self, e): if self.may_continue(): self.clear() self.close_event.emit() else: e.ignore() def resizeEvent(self, ev): self.canvas.adjustSize() self.canvas.update() super(DatabaseEditor, self).resizeEvent(ev)