import errno import os import YOSO from PyQt5.QtCore import (pyqtSlot, QIODevice, QModelIndex, QPointF, QRectF, QSaveFile, Qt) from PyQt5.QtGui import QColor, QPen, QPixmap, QTransform from PyQt5.QtWidgets import (QGraphicsItem, QGraphicsRectItem, QGraphicsScene, QGraphicsView) BBOX = QGraphicsItem.UserType + 1 def _mkRectF(p1, p2): rect = QRectF(p1.x(), p1.y(), p2.x() - p1.x(), p2.y() - p1.y()) return rect.normalized() class BoundingBoxItem(QGraphicsRectItem): _cls = None def __init__(self, cls, rect, model): super().__init__(rect) self._setClass(cls, model) self.setAcceptDrops(True) @property def number(self): return self._cls def type(self): return BBOX def _setClass(self, cls, model): self._cls = cls idx = model.findClass(self._cls) (hue, sat, val) = model.hsvF(self._cls) pen = QPen(QColor.fromHsvF(hue, sat, val, 0.6), 3) self.setPen(pen) if idx.isValid(): self.setToolTip(idx.data(Qt.DisplayRole)) else: self.setToolTip('{} - '.format(cls)) def dragEnterEvent(self, event): event.accept() def dropEvent(self, event): cls = int(event.mimeData().text()) self._setClass(cls, self.scene().model) self.scene().saveLabels() event.accept() class Scene(QGraphicsScene): _bbox_pen = None _bbox_start_pos = None _default_class = None _guide_pen = None _image = None _img_h = None _img_w = None _label_path = None _model = None _mouse_pos = None def __init__(self, parent): super().__init__(parent) self._guide_pen = QPen(Qt.black, 1) self.setDefaultClass(0) def _boxes(self): return list(filter(lambda i: i.type() == BBOX, self.items())) def saveLabels(self): if self._label_path != None: if len(self._boxes()) == 0: print('Removing empty "{}"'.format(self._label_path)) try: os.remove(self._label_path) except OSError as ex: if ex.errno != errno.ENOENT: raise else: lf = QSaveFile(self._label_path) if not lf.open(QIODevice.WriteOnly | QIODevice.Text): raise IOError('Cannot open "{}" for writing'.format(self._label_path)) for bbox in self._boxes(): rect = bbox.rect() c = bbox.number x = (rect.x() + rect.width() / 2) / self._img_w y = (rect.y() + rect.height() / 2) / self._img_h w = rect.width() / self._img_w h = rect.height() / self._img_h line = '{c:d} {x:.10f} {y:.10f} {w:.10f} {h:.10f}\n'.format(c=c, x=x, y=y, w=w, h=h) lf.write(line.encode()) if lf.commit(): print('Wrote "{}"'.format(self._label_path)) else: raise IOError('Cannot write "{}"'.format(self._label_path)) def _addBBox(self, p1, p2, cls): rect = _mkRectF(p1, p2).intersected(QRectF(self._image.rect())) square = abs(rect.width() * rect.height()) if square > 0.0: zvalue = 1.0 / square bbox = BoundingBoxItem(cls, rect, self._model) bbox.setZValue(zvalue) self.addItem(bbox) return True else: return False def _newBBox(self, p1, p2): if self._addBBox(p1, p2, self._default_class): self.invalidate() self.saveLabels() def setDefaultClass(self, cls): self._default_class = cls if self._model != None: (hue, sat, val) = self._model.hsvF(self._default_class) self._bbox_pen = QPen(QColor.fromHsvF(hue, sat, val), 2) def mouseMoveEvent(self, event): self._mouse_pos = event.scenePos() self.invalidate() super().mouseMoveEvent(event) def mousePressEvent(self, event): pos = event.scenePos() self._bbox_start_pos = pos def mouseReleaseEvent(self, event): if self._bbox_start_pos != None: self._newBBox(self._bbox_start_pos, event.scenePos()) self._bbox_start_pos = None self._mouse_pos = None def mouseDoubleClickEvent(self, event): item = self.itemAt(event.scenePos(), QTransform()) if item != None and item.type() == BBOX: self.removeItem(item) self.invalidate() self.saveLabels() def drawForeground(self, painter, rect): if self._mouse_pos != None: if self._bbox_start_pos == None: painter.setClipRect(rect) painter.setPen(self._guide_pen) painter.drawLine(self._mouse_pos.x(), rect.top(), self._mouse_pos.x(), rect.bottom()) painter.drawLine(rect.left(), self._mouse_pos.y(), rect.right(), self._mouse_pos.y()) else: painter.setPen(self._bbox_pen) painter.drawRect(_mkRectF(self._bbox_start_pos, self._mouse_pos)) def loadImage(self, image_path, label_path): self.clear() self._image = QPixmap(image_path) self._label_path = label_path self.setSceneRect(QRectF(self._image.rect())) self.addPixmap(self._image) lines = [] self._img_h = self._image.height() self._img_w = self._image.width() try: with open(self._label_path) as lbl: lines = lbl.readlines() except OSError as ex: if ex.errno != errno.ENOENT: raise for l in lines: m = YOSO.BBOX_RE.match(l) if m != None: cls = int(m.group('cls')) x = float(m.group('x')) y = float(m.group('y')) h = float(m.group('h')) w = float(m.group('w')) p1 = QPointF((x - w/2) * self._img_w, (y - h/2) * self._img_h) p2 = QPointF((x + w/2) * self._img_w, (y + h/2) * self._img_h) self._addBBox(p1, p2, cls) self.invalidate() @property def model(self): return self._model def setModel(self, model): self._model = model self.setDefaultClass(0) class Workspace(QGraphicsView): _scene = None def __init__(self): super().__init__() self._scene = Scene(self) self.setScene(self._scene) self.setMouseTracking(True) self.setEnabled(False) def _fit(self): scene_rect = self._scene.sceneRect() self.fitInView(scene_rect, Qt.KeepAspectRatio) @pyqtSlot(QModelIndex, QModelIndex) def setDefaultClass(self, curr, _): self._scene.setDefaultClass(curr.data(Qt.UserRole)) def resizeEvent(self, _): self._fit() def loadImage(self, image_path, label_path): self._scene.loadImage(image_path, label_path) self._fit() self.viewport().update() self.setEnabled(True) def setModel(self, model): self._scene.setModel(model)