diff options
| -rw-r--r-- | LICENSE | 13 | ||||
| -rw-r--r-- | README.rst | 92 | ||||
| -rw-r--r-- | YOSO/Classes.py | 88 | ||||
| -rw-r--r-- | YOSO/MainWindow.py | 190 | ||||
| -rw-r--r-- | YOSO/Workspace.py | 221 | ||||
| -rw-r--r-- | YOSO/__init__.py | 35 | ||||
| -rw-r--r-- | screenshots/yoso-roy-1.png | bin | 0 -> 213715 bytes | |||
| -rw-r--r-- | setup.py | 33 | ||||
| -rw-r--r-- | yoso.py | 6 |
9 files changed, 678 insertions, 0 deletions
@@ -0,0 +1,13 @@ + DO WHAT THE FUCK YOU WANT TO PUBLIC LICENSE + Version 2, December 2004 + + Copyright (C) 2004 Sam Hocevar <sam@hocevar.net> + + Everyone is permitted to copy and distribute verbatim or modified + copies of this license document, and changing it is allowed as long + as the name is changed. + + DO WHAT THE FUCK YOU WANT TO PUBLIC LICENSE + TERMS AND CONDITIONS FOR COPYING, DISTRIBUTION AND MODIFICATION + + 0. You just DO WHAT THE FUCK YOU WANT TO. diff --git a/README.rst b/README.rst new file mode 100644 index 0000000..f8a09a5 --- /dev/null +++ b/README.rst @@ -0,0 +1,92 @@ +You Only Show Once +================== + +YOSO is a GUI application providing a means of creating training data set for +`YOLO <http://pjreddie.com/darknet/yolo/>`_, a state-of-the-art, real-time object detection system. +It can be useful for other computer vision systems, because I believe that YOLO's +label file format is quite a good and flexible choice. + + +Requirements +============ + +- Python 3 +- PyQt5 + + +Data Directory Structure +======================== + +YOSO should be pointed to a data directory having this structure: + +:: + + data + |\_ classes + | \_ "0 - roy.jpg" + | \_ "1 - poli.jpg" + | \_ ... + |\_ images + | |\_ poli + | | \_ poli-001.jpg + | | \_ ... + | | \_ poli-009.jpg + | | \_ ... + | |\_ roy + | | \_ roy-1.jpg + | | \_ ... + | | \_ roy-99.jpg + | | \_ ... + | \_ ... + \_ labels + |\_ poli + | \_ poli-001.jpg.txt + | \_ ... + | \_ poli-009.jpg.txt + | \_ ... + |\_ roy + | \_ roy-1.jpg.txt + | \_ ... + | \_ roy-99.jpg.txt + | \_ ... + \_ ... + + +``data/classes`` must contain JPEG files describing object classes in format: ``<class number> - <short description>.jpg`` + +``data/images`` has arbitrary structure and contains JPEG or PNG images (``*.jpg``, or ``*.jpeg``, ``*.png``). +This is a training set of hundreds or thousands of images. + +``data/labels`` is managed by YOSO and has the same structure as ``data/images``. +All missed subdirectories will be created automatically. +Note that label files create by YOSO have different naming scheme, so you might have to update +`Darknet <http://pjreddie.com/darknet>`_ sources. + + +Controls +======== + +- To add a bounding box select a region with the mouse pointer. + Newly added bounding box will have object class currently selected + in the list on the right. + +- To delete a bounding box double click on it. + +- To change object class drag an item from the list on the right + and drop it into existing bounding box. + +Whenever a bounding box is added, deleted or changed, the result is automatically saved. + + +Screenshots +=========== + +.. image:: screenshots/yoso-roy-1.png + + +Video +===== + +https://www.youtube.com/watch?v=upPbaXq8wm0 + + 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()) + diff --git a/screenshots/yoso-roy-1.png b/screenshots/yoso-roy-1.png Binary files differnew file mode 100644 index 0000000..32456da --- /dev/null +++ b/screenshots/yoso-roy-1.png diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..85721e8 --- /dev/null +++ b/setup.py @@ -0,0 +1,33 @@ +from setuptools import setup, find_packages + +setup( + name = 'YOSO', + version = '0.1.0', + description = 'You Only Show Once', + long_description = 'A GUI tool to create training data for YOLO network by Darknet. See <http://pjreddie.com/darknet/yolo/>.', + author = 'Igor Pashev', + author_email = 'pashev.igor@gmail.com', + license = 'WTFPL', + data_files = [('', [ 'LICENSE' ])], + + classifiers = [ + 'Development Status :: 3 - Alpha', + 'Environment :: X11 Applications :: Qt', + 'Intended Audience :: Science/Research', + 'License :: Public Domain', + 'Programming Language :: Python :: 3.5', + 'Topic :: Multimedia :: Graphics :: Viewers', + 'Topic :: Scientific/Engineering :: Visualization', + 'Topic :: Utilities', + ], + + packages = find_packages(), + + install_requires = [ 'PyQt5' ], + + entry_points = { + 'console_scripts': [ + 'yoso=YOSO:main', + ], + }, +) @@ -0,0 +1,6 @@ +#!/usr/bin/env python3 + +from YOSO import main + +main() + |
