diff options
Diffstat (limited to 'YOSO')
| -rw-r--r-- | YOSO/Classes.py | 12 | ||||
| -rw-r--r-- | YOSO/MainWindow.py | 92 | ||||
| -rw-r--r-- | YOSO/Workspace.py | 60 | ||||
| -rw-r--r-- | YOSO/__init__.py | 21 |
4 files changed, 110 insertions, 75 deletions
diff --git a/YOSO/Classes.py b/YOSO/Classes.py index 9d5b09c..b48049d 100644 --- a/YOSO/Classes.py +++ b/YOSO/Classes.py @@ -1,6 +1,4 @@ -from PyQt5.QtCore import ( QAbstractListModel, QMimeData, - QModelIndex, Qt, QVariant ) - +from PyQt5.QtCore import QAbstractListModel, QMimeData, QModelIndex, Qt, QVariant class Class: @@ -27,7 +25,7 @@ class Class: @property def display(self): - return '{} - {}'.format(self._number, self._name) + return "{} - {}".format(self._number, self._name) class ClassListModel(QAbstractListModel): @@ -47,7 +45,7 @@ class ClassListModel(QAbstractListModel): return len(self._classes) def mimeTypes(self): - return 'text/plain' + return "text/plain" def mimeData(self, indeces): idx = indeces[0] @@ -78,11 +76,9 @@ class ClassListModel(QAbstractListModel): return self._classes def hsvF(self, num): - hue = 0.1 + 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 index bc54cc2..7978c77 100644 --- a/YOSO/MainWindow.py +++ b/YOSO/MainWindow.py @@ -4,14 +4,25 @@ 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 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 @@ -26,8 +37,12 @@ class MainWindow(QMainWindow): _workspace = None def openImages(self, directory): - image_it = QDirIterator(directory, YOSO.IMAGE_FILE_TEMPLATES, - QDir.Files, QDirIterator.Subdirectories) + 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()) @@ -43,7 +58,6 @@ class MainWindow(QMainWindow): self._image_spinner.setEnabled(True) self._image_spinner.setValue(1) - def openDataDir(self): dialog = QFileDialog(self) @@ -54,17 +68,19 @@ class MainWindow(QMainWindow): datadir = dialog.selectedFiles()[0] classesdir = YOSO.classesDir(datadir) - classesdir_it = QDirIterator(classesdir, YOSO.IMAGE_FILE_TEMPLATES, QDir.Files) + 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_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.sort(key=lambda c: c.number) classes_model = ClassListModel(classes) self._classes_view.setModel(classes_model) @@ -72,13 +88,18 @@ class MainWindow(QMainWindow): 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) + 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) + 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(): @@ -91,51 +112,51 @@ class MainWindow(QMainWindow): 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' + 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.setWindowTitle("YOSO - You Only Show Once") self.resize(800, 600) - self.move(qApp.desktop().availableGeometry().center() - self.frameGeometry().center()) - + 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 = 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 = 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 = 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.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 = 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.setShortcut("D") + self._next_image_action.setStatusTip("Show next image") self._next_image_action.triggered.connect(self.nextImage) menubar = self.menuBar() @@ -147,10 +168,9 @@ class MainWindow(QMainWindow): menubar.addSeparator() menubar.addAction(quit_action) - statusbar = self.statusBar() self._progress_bar = QProgressBar() - self._progress_bar.setFormat('%p% of %m') + self._progress_bar.setFormat("%p% of %m") self._image_spinner = QSpinBox() self._image_spinner.setEnabled(False) self._image_spinner.valueChanged.connect(self.loadImage) @@ -159,7 +179,6 @@ class MainWindow(QMainWindow): statusbar.addWidget(self._progress_bar) statusbar.show() - main_split = QSplitter(Qt.Horizontal) left_side = QWidget(main_split) @@ -182,10 +201,7 @@ class MainWindow(QMainWindow): 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 index 08057a9..e45ad55 100644 --- a/YOSO/Workspace.py +++ b/YOSO/Workspace.py @@ -3,18 +3,31 @@ import os import YOSO -from PyQt5.QtCore import (pyqtSlot, QIODevice, - QModelIndex, QPointF, QRectF, QSaveFile, Qt) +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 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 @@ -39,7 +52,7 @@ class BoundingBoxItem(QGraphicsRectItem): if idx.isValid(): self.setToolTip(idx.data(Qt.DisplayRole)) else: - self.setToolTip('{} - <unknown>'.format(cls)) + self.setToolTip("{} - <unknown>".format(cls)) def dragEnterEvent(self, event): event.accept() @@ -84,16 +97,20 @@ class Scene(QGraphicsScene): else: lf = QSaveFile(self._label_path) if not lf.open(QIODevice.WriteOnly | QIODevice.Text): - raise IOError('Cannot open "{}" for writing'.format(self._label_path)) + 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 + x = (rect.x() + rect.width() / 2) / self._img_w y = (rect.y() + rect.height() / 2) / self._img_h - w = rect.width() / self._img_w + 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) + 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(): @@ -151,8 +168,12 @@ class Scene(QGraphicsScene): 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()) + 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)) @@ -176,17 +197,16 @@ class Scene(QGraphicsScene): 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) + 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 @@ -196,7 +216,6 @@ class Scene(QGraphicsScene): self.setDefaultClass(0) - class Workspace(QGraphicsView): _scene = None @@ -227,4 +246,3 @@ class Workspace(QGraphicsView): def setModel(self, model): self._scene.setModel(model) - diff --git a/YOSO/__init__.py b/YOSO/__init__.py index 1decc61..6798d6f 100644 --- a/YOSO/__init__.py +++ b/YOSO/__init__.py @@ -5,24 +5,30 @@ import re import sys -IMAGE_FILE_TEMPLATES = ['*.png', '*.jpg', '*.jpeg'] +IMAGE_FILE_TEMPLATES = ["*.png", "*.jpg", "*.jpeg"] + def imagesDir(datadir): - return os.path.join(datadir, 'images') + return os.path.join(datadir, "images") + def labelsDir(datadir): - return os.path.join(datadir, 'labels') + return os.path.join(datadir, "labels") + def classesDir(datadir): - return os.path.join(datadir, 'classes') + 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>[^.]+).*$') +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]+')) + 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(): @@ -32,4 +38,3 @@ def main(): main_window.show() sys.exit(app.exec()) - |