# -*- coding: utf-8 -*- """ widgets.py -- A set of classes to extend widgets from pyqtgraph and pyqt for annotation purposes Copyright 2020 Monterey Bay Aquarium Research Institute Distributed under MIT license. See license.txt for more information. """ import os import cv2 import numpy as np from config import settings from pyqtgraph.Qt import QtCore, QtGui, QtWidgets class RectWidget(QtWidgets.QGraphicsWidget): rectHover = QtCore.Signal(object) def __init__(self, annotation, roi, index, parent=None, text_label='rect widget'): QtWidgets.QGraphicsWidget.__init__(self, parent) self.labelheight = 30 self.bordersize = 4 self.picdims = [240, 240] self.zoom = .5 self.text_label = text_label self.raw_image = roi self.raw_image = np.rot90(self.raw_image, 3, (0, 1)) self.pic = self.getpic(roi) self._boundingRect = QtCore.QRect() self.rectHeight = self.raw_image.shape[1] self.setAcceptHoverEvents(True) self.bgColor = QtCore.Qt.darkGray self.hoverColor = QtCore.Qt.lightGray self.isLastSelected = False self.isSelected = False self.isAnnotated = False self.forReview = False self.toDiscard = False self.annotation_object_index = index self.annotation = annotation self.annotation_path = annotation.xml_path def toqimage(self, img): height, width, bytesPerComponent = img.shape bytesPerLine = bytesPerComponent * width cv2.cvtColor(img, cv2.COLOR_BGR2RGB, img) qimg = QtGui.QImage(img.copy(), width, height, bytesPerLine, QtGui.QImage.Format_RGB888) return qimg def filename_to_box(self): coords = self.thumb_path.split('.')[-1].split('-') print(coords) self.ymax = int(coords[-1]) self.xmax = int(coords[-2]) self.ymin = int(coords[-3]) self.xmin = int(coords[-4]) def update_zoom(self, zoom): self.zoom = zoom self.boundingRect() self.updateGeometry() def get_image_path(self): ann = self.annotation data_path, filename = os.path.split(ann.xml_path) if self.forReview: image_path = os.path.join(data_path, settings.REVIEW_SUBDIR, ann.xml_tree.find('filename').text) elif self.toDiscard: image_path = os.path.join(data_path, settings.DISCARD_SUBDIR, ann.xml_tree.find('filename').text) else: image_path = os.path.join(data_path, ann.xml_tree.find('filename').text) return image_path def getFullImage(self): image_path = self.get_image_path() if not os.path.exists(image_path): return None full_img = cv2.imread(image_path) return np.rot90(full_img, 3, (0, 1)) def getBoxedImage(self): ann = self.annotation data_path, filename = os.path.split(ann.xml_path) image_path = self.get_image_path() if not os.path.exists(image_path): return None full_img = cv2.imread(image_path) for ind, ob in enumerate(self.annotation.xml_tree.findall('object')): if ind == self.annotation_object_index: color = (255, 255, 0) else: color = (75, 75, 75) xmin = int(ob.find('bndbox').find('xmin').text) xmax = int(ob.find('bndbox').find('xmax').text) ymin = int(ob.find('bndbox').find('ymin').text) ymax = int(ob.find('bndbox').find('ymax').text) if xmax - xmin <= 0 or ymax - ymin <= 0: print('bad box, not drawing...') else: cv2.rectangle(full_img, tuple([xmin, ymin]), tuple([xmax, ymax]), color, 2) return np.rot90(full_img, 3, (0, 1)) def boundingRect(self): #scale and zoom width = self.zoom*(self.picdims[0] + self.bordersize * 2) height = self.zoom*(self.picdims[1] + self.labelheight + self.bordersize * 2) thumb_widget_rect = QtCore.QRectF(0.0, 0.0, width, height) self._boundingRect = thumb_widget_rect return thumb_widget_rect def sizeHint(self, which, constraint=QtCore.QSizeF()): return self._boundingRect.size() def getpic(self, roi): height, width, channels = roi.shape if height >= width: scale = self.picdims[0]/height else: scale = self.picdims[0]/width new_width = int(width*scale) - 2*self.bordersize new_height = int(height*scale) - 2*self.bordersize roi = cv2.resize(roi, (new_width, new_height)) # center roi on dims w_pad = int((self.picdims[0] - new_width) / 2) h_pad = int((self.picdims[1] - new_height) / 2) roi = cv2.copyMakeBorder( roi, h_pad, h_pad, w_pad, w_pad, cv2.BORDER_CONSTANT, value=settings.BG_COLOR ) qimg = self.toqimage(roi) orpixmap = QtGui.QPixmap.fromImage(qimg) return orpixmap def paint(self, painter, option, widget): pen = QtGui.QPen() pen.setWidth(1) pen.setBrush(QtCore.Qt.black) painter.setPen(pen) # very simple selection and annotation logic if self.isSelected: fill_color = QtCore.Qt.green elif self.isAnnotated: fill_color = QtCore.Qt.yellow else: fill_color = QtCore.Qt.darkGray # fill behind image if self.isLastSelected: painter.fillRect(QtCore.QRect(0, 0, self.zoom*(self.pic.rect().width() + 2*self.bordersize), self.zoom*(self.pic.rect().height() + self.labelheight + 2*self.bordersize)), QtGui.QColor(61, 174, 233, 255)) # Fill label painter.fillRect(QtCore.QRect(self.zoom*self.bordersize, self.zoom*(self.bordersize + self.pic.rect().height()), self.zoom*self.pic.rect().width(), self.zoom*self.labelheight), fill_color) # Draw image painter.drawPixmap(QtCore.QRect(self.zoom*self.bordersize, self.zoom*self.bordersize, self.zoom*self.pic.rect().width(), self.zoom*self.pic.rect().height()), self.pic, self.pic.rect()) # Draw text text_rect = QtCore.QRect(0, self.zoom*(self.pic.rect().y() + self.pic.rect().height()), self.zoom*self.pic.rect().width(), self.zoom*self.labelheight) painter.drawText(text_rect, QtCore.Qt.AlignCenter, self.text_label) if self.toDiscard: painter.fillRect(QtCore.QRect(self.zoom * self.bordersize, self.zoom * (self.bordersize), self.zoom * self.pic.rect().width(), self.zoom * self.labelheight), QtCore.Qt.gray) text_rect = QtCore.QRect(0, self.zoom * (self.pic.rect().y()), self.zoom * self.pic.rect().width(), self.zoom * self.labelheight) painter.setPen(QtCore.Qt.red) painter.drawText(text_rect, QtCore.Qt.AlignCenter, "To Remove") if self.forReview: painter.fillRect(QtCore.QRect(self.zoom * self.bordersize, self.zoom * (self.bordersize), self.zoom * self.pic.rect().width(), self.zoom * self.labelheight), QtCore.Qt.gray) text_rect = QtCore.QRect(0, self.zoom * (self.pic.rect().y()), self.zoom * self.pic.rect().width(), self.zoom * self.labelheight) painter.setPen(QtCore.Qt.blue) painter.drawText(text_rect, QtCore.Qt.AlignCenter, "For Review") def mousePressEvent(self, event): self.isSelected = not self.isSelected self.update() self.rectHover.emit(self) def mouseReleaseEvent(self, event): pass def hoverEnterEvent(self, event): modifiers = QtWidgets.QApplication.keyboardModifiers() if modifiers == QtCore.Qt.ControlModifier: self.isSelected = not self.isSelected self.update() self.rectHover.emit(self)