Initial version 0.1.00.1.0

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())
+