aboutsummaryrefslogtreecommitdiff
path: root/YOSO
diff options
context:
space:
mode:
authorIgor Pashev <pashev.igor@gmail.com>2017-02-05 16:33:00 +0300
committerIgor Pashev <pashev.igor@gmail.com>2017-02-06 00:04:19 +0300
commit25469d39577aaa4e06a4b48ff7cdb5c2b4821b3c (patch)
treedde283030affdd066b2c4c490517d9770c4f13d6 /YOSO
downloadyoso-25469d39577aaa4e06a4b48ff7cdb5c2b4821b3c.tar.gz
Initial version 0.1.00.1.0
Diffstat (limited to 'YOSO')
-rw-r--r--YOSO/Classes.py88
-rw-r--r--YOSO/MainWindow.py190
-rw-r--r--YOSO/Workspace.py221
-rw-r--r--YOSO/__init__.py35
4 files changed, 534 insertions, 0 deletions
diff --git a/YOSO/Classes.py b/YOSO/Classes.py
new file mode 100644
index 0000000..9d5b09c
--- /dev/null
+++ b/YOSO/Classes.py
@@ -0,0 +1,88 @@
+from PyQt5.QtCore import ( QAbstractListModel, QMimeData,
+ QModelIndex, Qt, QVariant )
+
+
+
+class Class:
+ _image = None
+ _name = None
+ _number = None
+
+ def __init__(self, number, name, image):
+ self._image = image
+ self._name = name
+ self._number = number
+
+ @property
+ def name(self):
+ return self._name
+
+ @property
+ def number(self):
+ return self._number
+
+ @property
+ def image(self):
+ return self._image
+
+ @property
+ def display(self):
+ return '{} - {}'.format(self._number, self._name)
+
+
+class ClassListModel(QAbstractListModel):
+
+ _classes = None
+ _size = 0
+
+ def __init__(self, classes, parent=None):
+ super().__init__(parent)
+ self._classes = dict(enumerate(classes))
+ self._size = len(self._classes)
+
+ def flags(self, index):
+ return Qt.ItemIsSelectable | Qt.ItemIsDragEnabled | Qt.ItemIsEnabled
+
+ def rowCount(self, parent=QModelIndex()):
+ return len(self._classes)
+
+ def mimeTypes(self):
+ return 'text/plain'
+
+ def mimeData(self, indeces):
+ idx = indeces[0]
+ cls = idx.data(Qt.UserRole)
+ mime_data = QMimeData()
+ mime_data.setText(str(cls))
+ return mime_data
+
+ def data(self, index, role):
+ if index.isValid():
+ cl = self._classes[index.row()]
+ if role == Qt.DisplayRole:
+ return QVariant(cl.display)
+ if role == Qt.DecorationRole:
+ return QVariant(cl.image)
+ if role == Qt.UserRole:
+ return QVariant(cl.number)
+ return QVariant()
+
+ def findClass(self, num):
+ for i in range(self._size):
+ if self._classes[i].number == num:
+ return self.index(i, 0)
+ return QModelIndex()
+
+ @property
+ def classes(self):
+ return self._classes
+
+ def hsvF(self, num):
+ hue = 0.1
+ if self._size > 0:
+ hue += num / self._size
+ sat = 1.0
+ val = 1.0
+ return (hue, sat, val)
+
+
diff --git a/YOSO/MainWindow.py b/YOSO/MainWindow.py
new file mode 100644
index 0000000..5da8382
--- /dev/null
+++ b/YOSO/MainWindow.py
@@ -0,0 +1,190 @@
+from os import makedirs
+from os.path import basename
+import errno
+
+from PyQt5.QtCore import QDir, QDirIterator, QItemSelectionModel, Qt
+from PyQt5.QtGui import QPixmap
+from PyQt5.QtWidgets import ( QAction, QComboBox, QFileDialog,
+ QListView, QMainWindow, QProgressBar, QSpinBox, QSplitter, QVBoxLayout,
+ QWidget, qApp )
+
+from YOSO.Classes import Class, ClassListModel
+from YOSO.Workspace import Workspace
+import YOSO
+
+class MainWindow(QMainWindow):
+
+ _classes_view = None
+ _current_images = []
+ _image_dirs_combo_box = None
+ _image_spinner = None
+ _next_image_action = None
+ _prev_image_action = None
+ _progress_bar = None
+ _top_images_dir = None
+ _top_labels_dir = None
+ _workspace = None
+
+ def openImages(self, directory):
+ image_it = QDirIterator(directory, YOSO.IMAGE_FILE_TEMPLATES,
+ QDir.Files, QDirIterator.Subdirectories)
+ self._current_images = []
+ while image_it.hasNext():
+ self._current_images.append(image_it.next())
+ self._current_images.sort()
+
+ self._image_spinner.setRange(0, 0)
+ self._image_spinner.setEnabled(False)
+ self._image_spinner.setValue(0)
+ image_total = len(self._current_images)
+ if image_total > 0:
+ self._progress_bar.setRange(1, image_total + 1)
+ self._image_spinner.setRange(1, image_total)
+ self._image_spinner.setEnabled(True)
+ self._image_spinner.setValue(1)
+
+
+ def openDataDir(self):
+ dialog = QFileDialog(self)
+
+ dialog.setFileMode(QFileDialog.Directory)
+ dialog.setOption(QFileDialog.ReadOnly)
+
+ if dialog.exec():
+ datadir = dialog.selectedFiles()[0]
+
+ classesdir = YOSO.classesDir(datadir)
+ classesdir_it = QDirIterator(classesdir, YOSO.IMAGE_FILE_TEMPLATES, QDir.Files)
+ classes = []
+ while classesdir_it.hasNext():
+ class_image = classesdir_it.next()
+ match = YOSO.CLASSES_RE.match(basename(class_image))
+ if match != None:
+ class_num = int(match.group('cls'))
+ class_name = match.group('name')
+ class_object = Class(class_num, class_name, QPixmap(class_image))
+ classes.append(class_object)
+ classes.sort(key = lambda c: c.number)
+
+ classes_model = ClassListModel(classes)
+ self._classes_view.setModel(classes_model)
+ self._workspace.setModel(classes_model)
+ self._classes_view.setEnabled(len(classes) > 0)
+ selMod = self._classes_view.selectionModel()
+ selMod.currentChanged.connect(self._workspace.setDefaultClass)
+ selMod.setCurrentIndex(classes_model.index(0, 0), QItemSelectionModel.Select)
+
+ self._top_images_dir = YOSO.imagesDir(datadir)
+ self._top_labels_dir = YOSO.labelsDir(datadir)
+
+ imagedir_it = QDirIterator(self._top_images_dir,
+ QDir.AllDirs | QDir.NoDotAndDotDot, QDirIterator.Subdirectories)
+ self._image_dirs_combo_box.clear()
+ self._image_dirs_combo_box.addItem(self._top_images_dir)
+ while imagedir_it.hasNext():
+ img_dir = imagedir_it.next()
+ self._image_dirs_combo_box.addItem(img_dir)
+ lbl_dir = img_dir.replace(self._top_images_dir, self._top_labels_dir)
+ try:
+ makedirs(lbl_dir)
+ except OSError as ex:
+ if ex.errno != errno.EEXIST:
+ raise
+
+
+ def loadImage(self, i):
+ self._prev_image_action.setEnabled(i > 1)
+ self._next_image_action.setEnabled(i < self._image_spinner.maximum())
+ if 1 <= i <= self._image_spinner.maximum():
+ image_path = self._current_images[i - 1]
+ label_path = image_path.replace(self._top_images_dir, self._top_labels_dir) + '.txt'
+ self._workspace.loadImage(image_path, label_path)
+
+
+ def nextImage(self):
+ self._image_spinner.setValue(self._image_spinner.value() + 1)
+
+ def prevImage(self):
+ self._image_spinner.setValue(self._image_spinner.value() - 1)
+
+
+ def __init__(self):
+ super().__init__()
+
+ self.setWindowTitle('YOSO - You Only Show Once')
+ self.resize(800, 600)
+ self.move(qApp.desktop().availableGeometry().center() - self.frameGeometry().center())
+
+
+ quit_action = QAction('&Quit', self)
+ quit_action.setShortcut('Q')
+ quit_action.setStatusTip('Quit application')
+ quit_action.triggered.connect(qApp.quit)
+
+ open_action = QAction('&Open', self)
+ open_action.setShortcut('O')
+ open_action.setStatusTip('Open data directory')
+ open_action.triggered.connect(self.openDataDir)
+
+ self._prev_image_action = QAction('Prev (&A)', self)
+ self._prev_image_action.setEnabled(False)
+ self._prev_image_action.setShortcut('A')
+ self._prev_image_action.setStatusTip('Show previous image')
+ self._prev_image_action.triggered.connect(self.prevImage)
+
+ self._next_image_action = QAction('Next (&D)', self)
+ self._next_image_action.setEnabled(False)
+ self._next_image_action.setShortcut('D')
+ self._next_image_action.setStatusTip('Show next image')
+ self._next_image_action.triggered.connect(self.nextImage)
+
+ menubar = self.menuBar()
+ menubar.addAction(open_action)
+ menubar.addSeparator()
+ menubar.addAction(self._prev_image_action)
+ menubar.addAction(self._next_image_action)
+ menubar.addSeparator()
+ menubar.addAction(quit_action)
+
+
+ statusbar = self.statusBar()
+ self._progress_bar = QProgressBar()
+ self._progress_bar.setFormat('%p% of %m')
+ self._image_spinner = QSpinBox()
+ self._image_spinner.setEnabled(False)
+ self._image_spinner.valueChanged.connect(self.loadImage)
+ self._image_spinner.valueChanged.connect(self._progress_bar.setValue)
+ statusbar.addWidget(self._image_spinner)
+ statusbar.addWidget(self._progress_bar)
+ statusbar.show()
+
+
+ main_split = QSplitter(Qt.Horizontal)
+
+ left_side = QWidget(main_split)
+ right_side = QWidget(main_split)
+
+ right_layout = QVBoxLayout()
+ left_layout = QVBoxLayout()
+
+ right_side.setLayout(right_layout)
+ left_side.setLayout(left_layout)
+
+ self._workspace = Workspace()
+ self._classes_view = QListView()
+ self._classes_view.setEnabled(False)
+ self._classes_view.setDragEnabled(True)
+ right_layout.addWidget(self._classes_view)
+ left_layout.addWidget(self._workspace)
+
+ self._image_dirs_combo_box = QComboBox()
+ left_layout.addWidget(self._image_dirs_combo_box)
+ self._image_dirs_combo_box.currentTextChanged.connect(self.openImages)
+
+
+ main_split.setStretchFactor(0, 1)
+ main_split.setStretchFactor(1, 0)
+
+ self.setCentralWidget(main_split)
+
+
diff --git a/YOSO/Workspace.py b/YOSO/Workspace.py
new file mode 100644
index 0000000..79fe462
--- /dev/null
+++ b/YOSO/Workspace.py
@@ -0,0 +1,221 @@
+import errno
+
+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 )
+
+from pprint import pprint
+
+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('{} - <unknown>'.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 saveLabels(self):
+ if self._label_path != None:
+ lf = QSaveFile(self._label_path)
+ if not lf.open(QIODevice.WriteOnly | QIODevice.Text):
+ raise IOError('Cannot open "{}" for writing'.format(self._label_path))
+
+ for item in self.items():
+ if item.type() == BBOX:
+ rect = item.rect()
+ c = item.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, prev):
+ self._scene.setDefaultClass(curr.data(Qt.UserRole))
+
+ def resizeEvent(self, event):
+ 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)
+
diff --git a/YOSO/__init__.py b/YOSO/__init__.py
new file mode 100644
index 0000000..1decc61
--- /dev/null
+++ b/YOSO/__init__.py
@@ -0,0 +1,35 @@
+from PyQt5.QtWidgets import QApplication
+from YOSO.MainWindow import MainWindow
+import os
+import re
+import sys
+
+
+IMAGE_FILE_TEMPLATES = ['*.png', '*.jpg', '*.jpeg']
+
+def imagesDir(datadir):
+ return os.path.join(datadir, 'images')
+
+def labelsDir(datadir):
+ return os.path.join(datadir, 'labels')
+
+def classesDir(datadir):
+ return os.path.join(datadir, 'classes')
+
+# e. g. "012 - Midi skirt.jpg":
+CLASSES_RE = re.compile(r'^0*(?P<cls>\d+)\s*-\s*(?P<name>[^.]+).*$')
+
+# e. g. "1 0.57 0.42 0.17 0.6654"
+BBOX_RE = re.compile(
+ r'^\s*(?P<cls>\d+)\s+(?P<x>{float})\s+(?P<y>{float})\s+(?P<w>{float})\s+(?P<h>{float}).*$'.format(
+ float=r'([0-9]*[.])?[0-9]+'))
+
+
+def main():
+ app = QApplication(sys.argv)
+
+ main_window = MainWindow()
+ main_window.show()
+
+ sys.exit(app.exec())
+