aboutsummaryrefslogtreecommitdiff
path: root/YOSO
diff options
context:
space:
mode:
Diffstat (limited to 'YOSO')
-rw-r--r--YOSO/Classes.py12
-rw-r--r--YOSO/MainWindow.py92
-rw-r--r--YOSO/Workspace.py60
-rw-r--r--YOSO/__init__.py21
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())
-