# -*- coding: utf-8 -*- """ rectlabel.py -- Tools to implement a labeling UI for bounding boxes in images Copyright 2020 Monterey Bay Aquarium Research Institute Distributed under MIT license. See license.txt for more information. """ import time import numpy as np import pyqtgraph as pg from libs.logger import LOG from pyqtgraph.Qt import QtCore, QtGui, QtWidgets class BoundingBox(pg.RectROI): def __init__(self, view, pos, size=[1, 1], color=(255, 0, 0), label='ROI', all_labels=[]): pg.RectROI.__init__(self, pos, size, pen=pg.mkPen(color, width=3, style=QtCore.Qt.DashLine), invertible=True, removable=True, sideScalers=True) self.addScaleHandle([1, 0], [0, 1]) self.addScaleHandle([0, 1], [1, 0]) self.addTranslateHandle([0.5, 0.5]) self.sigRemoveRequested.connect(self.remove) self.sigRegionChanged.connect(self.draw_name) self.label = label self.all_labels = all_labels self.color = color self.textItem = None self.all_labels = all_labels self.view = view self.isNew = False self.isSelected = False self.removed = False self.view.addItem(self) self.draw_name() def update_label(self, label): self.label = label self.draw_name() def remove(self, dummy): self.view.removeItem(self.textItem) self.view.removeItem(self) self.removed = True def get_box(self): x, y = self.pos() w, h = self.size() # need to convert box coordinates to account for positive or # negative width/height values from pyqtgraph roi object if h > 0: y = y + h if w < 0: x = x + w return [x, y, np.abs(w), np.abs(h)] def draw_name(self): roiSize = self.size() w = roiSize[0] h = roiSize[1] if w < 0: x = self.pos().x() - np.abs(w) else: x = self.pos().x() if h < 0: y = self.pos().y() - np.abs(h) else: y = self.pos().y() if self.textItem is None: self.textItem = pg.TextItem(text=self.label, color=np.array(self.color)/2, fill=(200, 200, 200)) self.textItem.setPos(x, y) self.view.addItem(self.textItem) else: self.textItem.setText(self.label) self.textItem.setPos(x, y) # On right-click, raise the context menu def mouseClickEvent(self, ev): if ev.button() == QtCore.Qt.RightButton: if self.raiseContextMenu(ev): ev.accept() def raiseContextMenu(self, ev): menu = self.getContextMenus() # Let the scene add on to the end of our context menu # (this is optional) #menu = self.scene().addParentContextMenus(self, menu, ev) pos = ev.screenPos() menu.popup(QtCore.QPoint(pos.x(), pos.y())) return True # This method will be called when this item's _children_ want to raise # a context menu that includes their parents' menus. def getContextMenus(self, event=None): if self.menu is None: self.menu = QtGui.QMenu() self.menu.setTitle(self.label + " Select Label") remove = QtGui.QAction("Remove ROI", self.menu) remove.triggered.connect(self.remove) self.menu.addAction(remove) self.menu.remove = remove self.menu.addSeparator() class_label = QtGui.QAction("Class Label", self.menu) self.menu.addAction(class_label) select_label = QtWidgets.QWidgetAction(self.menu) cb = QtWidgets.QComboBox() cb.setStyleSheet("") for l in self.all_labels: cb.addItem(l) # Set the current selection to the current label for ind, l in enumerate(self.all_labels): if self.all_labels[ind].lower().strip(' ') == self.label.lower().strip(' '): cb.setCurrentIndex(ind) break cb.currentTextChanged.connect(self.update_label) select_label.setDefaultWidget(cb) self.menu.addAction(select_label) self.menu.select_label = select_label self.menu.cb = cb return self.menu class RectLabel: def __init__(self, graphics_view, annotation=None, all_labels=[]): self.vb = pg.ViewBox() self.graphics_view = graphics_view self.graphics_view.setCentralItem(self.vb) self.vb.setAspectLocked() self.roiDetail = pg.ImageItem() self.vb.addItem(self.roiDetail) self.annotation = annotation self.all_labels = all_labels self.rois = [] self.dragging = False self.label = "" self.graphics_view.scene().sigMouseClicked.connect(self.onMouseClick) self.graphics_view.scene().sigMouseMoved.connect(self.onMouseMove) def apply_label(self, label): for r in self.rois: if r.isSelected: r.label = label r.draw_name() def set_label(self, label): self.label = label # in most cases I don't think we want to change existing roi labels # this way # Added in as apply_label method above per discussion with Alex #for r in self.rois: # if r.isSelected: # r.label = label # r.draw_name() def add_annotation(self, obj_index, rect): for ind, ob in enumerate(rect.annotation.xml_tree.findall('object')): if ind == obj_index: color = (52, 161, 235) else: color = (143, 52, 235) xmin, ymin, xmax, ymax = rect.annotation.get_box(ind) height = int(rect.annotation.xml_tree.find('size').find('height').text) if xmax - xmin <= 0 or ymax - ymin <= 0: LOG.warn('bad box, not adding') else: # note invert y-axis for bb from rectlabel bb = BoundingBox(self.vb, [xmin, height-ymin], [xmax-xmin, -1*(ymax-ymin)], color=color, label=ob.find('name').text, all_labels=self.all_labels) bb.isSelected = rect.isSelected self.rois.append(bb) self.annotation = rect.annotation def save_all(self): # append any new boxes to the xml height = int(self.annotation.xml_tree.find('size').find('height').text) objs = self.annotation.xml_tree.findall('object') # clear old objects self.annotation.remove_boxes() for ind, roi in enumerate(self.rois): if roi.removed: continue box = roi.get_box() xmin = box[0] ymin = height - box[1] xmax = box[0] + box[2] ymax = ymin + np.abs(box[3]) self.annotation.add_box([(xmin, ymin), (xmax, ymax), int(time.time()), roi.label, 1.0]) self.annotation.save_xml() def mapToItem(self, pos): pt = self.vb.mapSceneToView(pos) pt = self.vb.mapFromViewToItem(self.roiDetail, pt) return pt def clear(self): for roi in self.rois: roi.remove([]) self.rois = [] def onMouseClick(self, event): mousePoint = self.mapToItem(event.scenePos()) if QtCore.Qt.ControlModifier == event.modifiers(): self.dragging = True bb = BoundingBox(self.vb, [mousePoint.x(), mousePoint.y()], [2, 2], color=(200, 200, 50), label=self.label, all_labels=self.all_labels) bb.isNew = True self.rois.append(bb) def onMouseMove(self, event): mousePoint = self.mapToItem(event) if self.dragging and QtCore.Qt.ControlModifier == QtWidgets.QApplication.keyboardModifiers(): width = -1 * (self.rois[-1].pos().x() - mousePoint.x()) height = -1 * (self.rois[-1].pos().y() - mousePoint.y()) self.rois[-1].setSize([width, height]) else: self.dragging = False