diff options
Diffstat (limited to 'YOSO/MainWindow.py')
| -rw-r--r-- | YOSO/MainWindow.py | 190 |
1 files changed, 190 insertions, 0 deletions
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) + + |