labelImg.py 52.2 KB
Newer Older
TzuTa Lin's avatar
TzuTa Lin committed
1
#!/usr/bin/env python
Jiye Qian's avatar
Jiye Qian committed
2
# -*- coding: utf-8 -*-
3
import codecs
TzuTa Lin's avatar
TzuTa Lin committed
4
5
6
7
8
9
10
11
import os.path
import re
import sys
import subprocess

from functools import partial
from collections import defaultdict

12
13
14
15
16
17
try:
    from PyQt5.QtGui import *
    from PyQt5.QtCore import *
    from PyQt5.QtWidgets import *
except ImportError:
    # needed for py3+qt4
18
19
    # Ref:
    # http://pyqt.sourceforge.net/Docs/PyQt4/incompatible_apis.html
20
    # http://stackoverflow.com/questions/21217399/pyqt4-qtcore-qvariant-object-instead-of-a-string
21
22
23
24
25
    if sys.version_info.major >= 3:
        import sip
        sip.setapi('QVariant', 2)
    from PyQt4.QtGui import *
    from PyQt4.QtCore import *
TzuTa Lin's avatar
TzuTa Lin committed
26
27

import resources
tzutalin's avatar
tzutalin committed
28
# Add internal libs
29
from libs.constants import *
30
from libs.lib import struct, newAction, newIcon, addActions, fmtShortcut, generateColorByText
tzutalin's avatar
tzutalin committed
31
from libs.settings import Settings
32
33
34
35
36
37
38
39
40
41
from libs.shape import Shape, DEFAULT_LINE_COLOR, DEFAULT_FILL_COLOR
from libs.canvas import Canvas
from libs.zoomWidget import ZoomWidget
from libs.labelDialog import LabelDialog
from libs.colorDialog import ColorDialog
from libs.labelFile import LabelFile, LabelFileError
from libs.toolBar import ToolBar
from libs.pascal_voc_io import PascalVocReader
from libs.pascal_voc_io import XML_EXT
from libs.ustr import ustr
tzutalin's avatar
tzutalin committed
42
from libs.version import __version__
TzuTa Lin's avatar
TzuTa Lin committed
43
44
45

__appname__ = 'labelImg'

46
47
# Utility functions and classes.

48
49
50
51
52
53
54
55
def have_qstring():
    '''p3/qt5 get rid of QString wrapper as py3 has native unicode str type'''
    return not (sys.version_info.major >= 3 or QT_VERSION_STR.startswith('5.'))

def util_qt_strlistclass():
    return QStringList if have_qstring() else list


TzuTa Lin's avatar
TzuTa Lin committed
56
class WindowMixin(object):
57

TzuTa Lin's avatar
TzuTa Lin committed
58
59
60
61
62
63
64
65
66
    def menu(self, title, actions=None):
        menu = self.menuBar().addMenu(title)
        if actions:
            addActions(menu, actions)
        return menu

    def toolbar(self, title, actions=None):
        toolbar = ToolBar(title)
        toolbar.setObjectName(u'%sToolBar' % title)
67
        # toolbar.setOrientation(Qt.Vertical)
TzuTa Lin's avatar
TzuTa Lin committed
68
69
70
71
72
73
74
        toolbar.setToolButtonStyle(Qt.ToolButtonTextUnderIcon)
        if actions:
            addActions(toolbar, actions)
        self.addToolBar(Qt.LeftToolBarArea, toolbar)
        return toolbar


75
76
# PyQt5: TypeError: unhashable type: 'QListWidgetItem'
class HashableQListWidgetItem(QListWidgetItem):
77

78
79
    def __init__(self, *args):
        super(HashableQListWidgetItem, self).__init__(*args)
80

81
82
83
84
    def __hash__(self):
        return hash(id(self))


TzuTa Lin's avatar
TzuTa Lin committed
85
class MainWindow(QMainWindow, WindowMixin):
86
    FIT_WINDOW, FIT_WIDTH, MANUAL_ZOOM = list(range(3))
TzuTa Lin's avatar
TzuTa Lin committed
87

88
    def __init__(self, defaultFilename=None, defaultPrefdefClassFile=None):
TzuTa Lin's avatar
TzuTa Lin committed
89
90
        super(MainWindow, self).__init__()
        self.setWindowTitle(__appname__)
91
92
93
94
95
96

        # Load setting in the main thread
        self.settings = Settings()
        self.settings.load()
        settings = self.settings

TzuTa Lin's avatar
Typo    
TzuTa Lin committed
97
        # Save as Pascal voc xml
TzuTa Lin's avatar
TzuTa Lin committed
98
        self.defaultSaveDir = None
TzuTa Lin's avatar
Typo    
TzuTa Lin committed
99
        self.usingPascalVocFormat = True
TzuTa Lin's avatar
TzuTa Lin committed
100
101
102
        # For loading all image under a directory
        self.mImgList = []
        self.dirname = None
103
        self.labelHist = []
104
        self.lastOpenDir = None
TzuTa Lin's avatar
TzuTa Lin committed
105
106
107
108
109
110
111

        # Whether we need to save or not.
        self.dirty = False

        self._noSelectionSlot = False
        self._beginner = True
        self.screencastViewer = "firefox"
TzuTa Lin's avatar
TzuTa Lin committed
112
        self.screencast = "https://youtu.be/p0nR2YsCY_U"
TzuTa Lin's avatar
TzuTa Lin committed
113

114
115
116
        # Load predefined classes to the list
        self.loadPredefinedClasses(defaultPrefdefClassFile)

TzuTa Lin's avatar
TzuTa Lin committed
117
        # Main widgets and related state.
118
        self.labelDialog = LabelDialog(parent=self, listItem=self.labelHist)
119

TzuTa Lin's avatar
TzuTa Lin committed
120
121
        self.itemsToShapes = {}
        self.shapesToItems = {}
122
        self.prevLabelText = ''
TzuTa Lin's avatar
TzuTa Lin committed
123
124
125

        listLayout = QVBoxLayout()
        listLayout.setContentsMargins(0, 0, 0, 0)
126

tzutalin's avatar
tzutalin committed
127
        # Create a widget for using default label
128
129
        self.useDefaultLabelCheckbox = QCheckBox(u'Use default label')
        self.useDefaultLabelCheckbox.setChecked(False)
tzutalin's avatar
tzutalin committed
130
        self.defaultLabelTextLine = QLineEdit()
131
132
133
134
135
        useDefaultLabelQHBoxLayout = QHBoxLayout()
        useDefaultLabelQHBoxLayout.addWidget(self.useDefaultLabelCheckbox)
        useDefaultLabelQHBoxLayout.addWidget(self.defaultLabelTextLine)
        useDefaultLabelContainer = QWidget()
        useDefaultLabelContainer.setLayout(useDefaultLabelQHBoxLayout)
tzutalin's avatar
tzutalin committed
136
137
138

        # Create a widget for edit and diffc button
        self.diffcButton = QCheckBox(u'difficult')
139
        self.diffcButton.setChecked(False)
140
        self.diffcButton.stateChanged.connect(self.btnstate)
tzutalin's avatar
tzutalin committed
141
142
        self.editButton = QToolButton()
        self.editButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
143

144
        # Add some of widgets to listLayout
tzutalin's avatar
tzutalin committed
145
        listLayout.addWidget(self.editButton)
146
        listLayout.addWidget(self.diffcButton)
147
        listLayout.addWidget(useDefaultLabelContainer)
148

tzutalin's avatar
tzutalin committed
149
150
151
152
153
154
155
156
157
158
        # Create and add a widget for showing current label items
        self.labelList = QListWidget()
        labelListContainer = QWidget()
        labelListContainer.setLayout(listLayout)
        self.labelList.itemActivated.connect(self.labelSelectionChanged)
        self.labelList.itemSelectionChanged.connect(self.labelSelectionChanged)
        self.labelList.itemDoubleClicked.connect(self.editLabel)
        # Connect to itemChanged to detect checkbox changes.
        self.labelList.itemChanged.connect(self.labelItemChanged)
        listLayout.addWidget(self.labelList)
TzuTa Lin's avatar
TzuTa Lin committed
159
160
161

        self.dock = QDockWidget(u'Box Labels', self)
        self.dock.setObjectName(u'Labels')
tzutalin's avatar
tzutalin committed
162
        self.dock.setWidget(labelListContainer)
TzuTa Lin's avatar
TzuTa Lin committed
163

164
165
        # Tzutalin 20160906 : Add file list and dock to move faster
        self.fileListWidget = QListWidget()
tzutalin's avatar
tzutalin committed
166
        self.fileListWidget.itemDoubleClicked.connect(self.fileitemDoubleClicked)
167
168
169
        filelistLayout = QVBoxLayout()
        filelistLayout.setContentsMargins(0, 0, 0, 0)
        filelistLayout.addWidget(self.fileListWidget)
tzutalin's avatar
tzutalin committed
170
171
        fileListContainer = QWidget()
        fileListContainer.setLayout(filelistLayout)
172
173
        self.filedock = QDockWidget(u'File List', self)
        self.filedock.setObjectName(u'Files')
tzutalin's avatar
tzutalin committed
174
        self.filedock.setWidget(fileListContainer)
175

TzuTa Lin's avatar
TzuTa Lin committed
176
177
178
        self.zoomWidget = ZoomWidget()
        self.colorDialog = ColorDialog(parent=self)

179
        self.canvas = Canvas(parent=self)
TzuTa Lin's avatar
TzuTa Lin committed
180
181
182
183
184
185
186
187
        self.canvas.zoomRequest.connect(self.zoomRequest)

        scroll = QScrollArea()
        scroll.setWidget(self.canvas)
        scroll.setWidgetResizable(True)
        self.scrollBars = {
            Qt.Vertical: scroll.verticalScrollBar(),
            Qt.Horizontal: scroll.horizontalScrollBar()
188
        }
Lars Klein's avatar
Lars Klein committed
189
        self.scrollArea = scroll
TzuTa Lin's avatar
TzuTa Lin committed
190
191
192
193
194
195
196
197
198
        self.canvas.scrollRequest.connect(self.scrollRequest)

        self.canvas.newShape.connect(self.newShape)
        self.canvas.shapeMoved.connect(self.setDirty)
        self.canvas.selectionChanged.connect(self.shapeSelectionChanged)
        self.canvas.drawingPolygon.connect(self.toggleDrawingSensitive)

        self.setCentralWidget(scroll)
        self.addDockWidget(Qt.RightDockWidgetArea, self.dock)
199
200
        # Tzutalin 20160906 : Add file list and dock to move faster
        self.addDockWidget(Qt.RightDockWidgetArea, self.filedock)
tzutalin's avatar
tzutalin committed
201
        self.filedock.setFeatures(QDockWidget.DockWidgetFloatable)
202

tzutalin's avatar
tzutalin committed
203
        self.dockFeatures = QDockWidget.DockWidgetClosable | QDockWidget.DockWidgetFloatable
TzuTa Lin's avatar
TzuTa Lin committed
204
205
206
207
208
        self.dock.setFeatures(self.dock.features() ^ self.dockFeatures)

        # Actions
        action = partial(newAction, self)
        quit = action('&Quit', self.close,
209
                      'Ctrl+Q', 'quit', u'Quit application')
210

TzuTa Lin's avatar
TzuTa Lin committed
211
        open = action('&Open', self.openFile,
212
                      'Ctrl+O', 'open', u'Open image or label file')
TzuTa Lin's avatar
TzuTa Lin committed
213

214
        opendir = action('&Open Dir', self.openDirDialog,
215
                         'Ctrl+u', 'open', u'Open Dir')
TzuTa Lin's avatar
TzuTa Lin committed
216

217
        changeSavedir = action('&Change Save Dir', self.changeSavedirDialog,
218
                               'Ctrl+r', 'open', u'Change default saved Annotation dir')
TzuTa Lin's avatar
TzuTa Lin committed
219

220
        openAnnotation = action('&Open Annotation', self.openAnnotationDialog,
221
                                'Ctrl+Shift+O', 'open', u'Open Annotation')
222

TzuTa Lin's avatar
TzuTa Lin committed
223
        openNextImg = action('&Next Image', self.openNextImg,
224
                             'd', 'next', u'Open Next')
TzuTa Lin's avatar
TzuTa Lin committed
225

tzutalin's avatar
tzutalin committed
226
        openPrevImg = action('&Prev Image', self.openPrevImg,
227
                             'a', 'prev', u'Open Prev')
tzutalin's avatar
tzutalin committed
228

Thibaut Mattio's avatar
Thibaut Mattio committed
229
230
231
        verify = action('&Verify Image', self.verifyImg,
                        'space', 'verify', u'Verify Image')

TzuTa Lin's avatar
TzuTa Lin committed
232
        save = action('&Save', self.saveFile,
233
                      'Ctrl+S', 'save', u'Save labels to file', enabled=False)
tzutalin's avatar
tzutalin committed
234

TzuTa Lin's avatar
TzuTa Lin committed
235
        saveAs = action('&Save As', self.saveFileAs,
tzutalin's avatar
tzutalin committed
236
237
238
                        'Ctrl+Shift+S', 'save-as', u'Save labels to a different file', enabled=False)

        close = action('&Close', self.closeFile, 'Ctrl+W', 'close', u'Close current file')
zhangjie's avatar
zhangjie committed
239

tzutalin's avatar
tzutalin committed
240
241
        resetAll = action('&ResetAll', self.resetAll, None, 'resetall', u'Reset all')

242
        color1 = action('Box Line Color', self.chooseColor1,
243
                        'Ctrl+L', 'color_line', u'Choose Box line color')
TzuTa Lin's avatar
TzuTa Lin committed
244
245

        createMode = action('Create\nRectBox', self.setCreateMode,
246
                            'w', 'new', u'Start drawing Boxs', enabled=False)
TzuTa Lin's avatar
TzuTa Lin committed
247
        editMode = action('&Edit\nRectBox', self.setEditMode,
248
                          'Ctrl+J', 'edit', u'Move and edit Boxs', enabled=False)
TzuTa Lin's avatar
TzuTa Lin committed
249
250

        create = action('Create\nRectBox', self.createShape,
251
                        'w', 'new', u'Draw a new Box', enabled=False)
TzuTa Lin's avatar
TzuTa Lin committed
252
        delete = action('Delete\nRectBox', self.deleteSelectedShape,
253
                        'Delete', 'delete', u'Delete', enabled=False)
TzuTa Lin's avatar
TzuTa Lin committed
254
        copy = action('&Duplicate\nRectBox', self.copySelectedShape,
255
256
                      'Ctrl+D', 'copy', u'Create a duplicate of the selected Box',
                      enabled=False)
TzuTa Lin's avatar
TzuTa Lin committed
257
258

        advancedMode = action('&Advanced Mode', self.toggleAdvancedMode,
259
260
                              'Ctrl+Shift+A', 'expert', u'Switch to advanced mode',
                              checkable=True)
TzuTa Lin's avatar
TzuTa Lin committed
261
262

        hideAll = action('&Hide\nRectBox', partial(self.togglePolygons, False),
263
264
                         'Ctrl+H', 'hide', u'Hide all Boxs',
                         enabled=False)
TzuTa Lin's avatar
TzuTa Lin committed
265
        showAll = action('&Show\nRectBox', partial(self.togglePolygons, True),
266
267
                         'Ctrl+A', 'hide', u'Show all Boxs',
                         enabled=False)
TzuTa Lin's avatar
TzuTa Lin committed
268

tzutalin's avatar
tzutalin committed
269
270
        help = action('&Tutorial', self.showTutorialDialog, None, 'help', u'Show demos')
        showInfo = action('&Information', self.showInfoDialog, None, 'help', u'Information')
TzuTa Lin's avatar
TzuTa Lin committed
271
272
273
274

        zoom = QWidgetAction(self)
        zoom.setDefaultWidget(self.zoomWidget)
        self.zoomWidget.setWhatsThis(
275
276
277
            u"Zoom in or out of the image. Also accessible with"
            " %s and %s from the canvas." % (fmtShortcut("Ctrl+[-+]"),
                                             fmtShortcut("Ctrl+Wheel")))
TzuTa Lin's avatar
TzuTa Lin committed
278
279
280
        self.zoomWidget.setEnabled(False)

        zoomIn = action('Zoom &In', partial(self.addZoom, 10),
281
                        'Ctrl++', 'zoom-in', u'Increase zoom level', enabled=False)
TzuTa Lin's avatar
TzuTa Lin committed
282
        zoomOut = action('&Zoom Out', partial(self.addZoom, -10),
283
                         'Ctrl+-', 'zoom-out', u'Decrease zoom level', enabled=False)
TzuTa Lin's avatar
TzuTa Lin committed
284
        zoomOrg = action('&Original size', partial(self.setZoom, 100),
285
                         'Ctrl+=', 'zoom', u'Zoom to original size', enabled=False)
TzuTa Lin's avatar
TzuTa Lin committed
286
        fitWindow = action('&Fit Window', self.setFitWindow,
287
288
                           'Ctrl+F', 'fit-window', u'Zoom follows window size',
                           checkable=True, enabled=False)
TzuTa Lin's avatar
TzuTa Lin committed
289
        fitWidth = action('Fit &Width', self.setFitWidth,
290
291
                          'Ctrl+Shift+F', 'fit-width', u'Zoom follows window width',
                          checkable=True, enabled=False)
TzuTa Lin's avatar
TzuTa Lin committed
292
        # Group zoom controls into a list for easier toggling.
293
294
        zoomActions = (self.zoomWidget, zoomIn, zoomOut,
                       zoomOrg, fitWindow, fitWidth)
TzuTa Lin's avatar
TzuTa Lin committed
295
296
297
298
299
300
301
302
303
        self.zoomMode = self.MANUAL_ZOOM
        self.scalers = {
            self.FIT_WINDOW: self.scaleFitWindow,
            self.FIT_WIDTH: self.scaleFitWidth,
            # Set to one to scale to 100% when loading files.
            self.MANUAL_ZOOM: lambda: 1,
        }

        edit = action('&Edit Label', self.editLabel,
304
305
                      'Ctrl+E', 'edit', u'Modify the label of the selected Box',
                      enabled=False)
TzuTa Lin's avatar
TzuTa Lin committed
306
307
308
        self.editButton.setDefaultAction(edit)

        shapeLineColor = action('Shape &Line Color', self.chshapeLineColor,
309
310
                                icon='color_line', tip=u'Change the line color for this specific shape',
                                enabled=False)
TzuTa Lin's avatar
TzuTa Lin committed
311
        shapeFillColor = action('Shape &Fill Color', self.chshapeFillColor,
312
313
                                icon='color', tip=u'Change the fill color for this specific shape',
                                enabled=False)
TzuTa Lin's avatar
TzuTa Lin committed
314
315
316
317
318
319
320
321
322

        labels = self.dock.toggleViewAction()
        labels.setText('Show/Hide Label Panel')
        labels.setShortcut('Ctrl+Shift+L')

        # Lavel list context menu.
        labelMenu = QMenu()
        addActions(labelMenu, (edit, delete))
        self.labelList.setContextMenuPolicy(Qt.CustomContextMenu)
323
324
        self.labelList.customContextMenuRequested.connect(
            self.popLabelListMenu)
TzuTa Lin's avatar
TzuTa Lin committed
325
326

        # Store actions for further handling.
tzutalin's avatar
tzutalin committed
327
        self.actions = struct(save=save, saveAs=saveAs, open=open, close=close, resetAll = resetAll,
328
                              lineColor=color1, create=create, delete=delete, edit=edit, copy=copy,
329
330
331
332
333
334
                              createMode=createMode, editMode=editMode, advancedMode=advancedMode,
                              shapeLineColor=shapeLineColor, shapeFillColor=shapeFillColor,
                              zoom=zoom, zoomIn=zoomIn, zoomOut=zoomOut, zoomOrg=zoomOrg,
                              fitWindow=fitWindow, fitWidth=fitWidth,
                              zoomActions=zoomActions,
                              fileMenuActions=(
tzutalin's avatar
tzutalin committed
335
                                  open, opendir, save, saveAs, close, resetAll, quit),
336
337
                              beginner=(), advanced=(),
                              editMenu=(edit, copy, delete,
338
                                        None, color1),
339
340
341
342
343
344
                              beginnerContext=(create, edit, copy, delete),
                              advancedContext=(createMode, editMode, edit, copy,
                                               delete, shapeLineColor, shapeFillColor),
                              onLoadActive=(
                                  close, create, createMode, editMode),
                              onShapesPresent=(saveAs, hideAll, showAll))
TzuTa Lin's avatar
TzuTa Lin committed
345
346

        self.menus = struct(
347
348
349
350
351
352
            file=self.menu('&File'),
            edit=self.menu('&Edit'),
            view=self.menu('&View'),
            help=self.menu('&Help'),
            recentFiles=QMenu('Open &Recent'),
            labelList=labelMenu)
TzuTa Lin's avatar
TzuTa Lin committed
353

354
        # Auto saving : Enable auto saving if pressing next
355
356
        self.autoSaving = QAction("Auto Saving", self)
        self.autoSaving.setCheckable(True)
357
        self.autoSaving.setChecked(settings.get(SETTING_AUTO_SAVE, False))
358
359
360
361
        # Sync single class mode from PR#106
        self.singleClassMode = QAction("Single Class Mode", self)
        self.singleClassMode.setShortcut("Ctrl+Shift+S")
        self.singleClassMode.setCheckable(True)
362
        self.singleClassMode.setChecked(settings.get(SETTING_SINGLE_CLASS, False))
363
364
        self.lastLabel = None

TzuTa Lin's avatar
TzuTa Lin committed
365
        addActions(self.menus.file,
tzutalin's avatar
tzutalin committed
366
                   (open, opendir, changeSavedir, openAnnotation, self.menus.recentFiles, save, saveAs, close, resetAll, quit))
tzutalin's avatar
tzutalin committed
367
        addActions(self.menus.help, (help, showInfo))
TzuTa Lin's avatar
TzuTa Lin committed
368
        addActions(self.menus.view, (
369
370
            self.autoSaving,
            self.singleClassMode,
TzuTa Lin's avatar
TzuTa Lin committed
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
            labels, advancedMode, None,
            hideAll, showAll, None,
            zoomIn, zoomOut, zoomOrg, None,
            fitWindow, fitWidth))

        self.menus.file.aboutToShow.connect(self.updateFileMenu)

        # Custom context menu for the canvas widget:
        addActions(self.canvas.menus[0], self.actions.beginnerContext)
        addActions(self.canvas.menus[1], (
            action('&Copy here', self.copyShape),
            action('&Move here', self.moveShape)))

        self.tools = self.toolbar('Tools')
        self.actions.beginner = (
386
            open, opendir, changeSavedir, openNextImg, openPrevImg, verify, save, None, create, copy, delete, None,
TzuTa Lin's avatar
TzuTa Lin committed
387
388
389
            zoomIn, zoom, zoomOut, fitWindow, fitWidth)

        self.actions.advanced = (
390
            open, opendir, changeSavedir, openNextImg, openPrevImg, save, None,
TzuTa Lin's avatar
TzuTa Lin committed
391
392
393
394
395
396
397
398
            createMode, editMode, None,
            hideAll, showAll)

        self.statusBar().showMessage('%s started.' % __appname__)
        self.statusBar().show()

        # Application state.
        self.image = QImage()
399
        self.filePath = ustr(defaultFilename)
TzuTa Lin's avatar
TzuTa Lin committed
400
401
402
403
404
405
        self.recentFiles = []
        self.maxRecent = 7
        self.lineColor = None
        self.fillColor = None
        self.zoom_level = 100
        self.fit_window = False
406
        # Add Chris
407
        self.difficult = False
TzuTa Lin's avatar
TzuTa Lin committed
408

409
410
411
412
413
414
415
416
417
418
        ## Fix the compatible issue for qt4 and qt5. Convert the QStringList to python list
        if settings.get(SETTING_RECENT_FILES):
            if have_qstring():
                recentFileQStringList = settings.get(SETTING_RECENT_FILES)
                self.recentFiles = [ustr(i) for i in recentFileQStringList]
            else:
                self.recentFiles = recentFileQStringList = settings.get(SETTING_RECENT_FILES)

        size = settings.get(SETTING_WIN_SIZE, QSize(600, 500))
        position = settings.get(SETTING_WIN_POSE, QPoint(0, 0))
TzuTa Lin's avatar
TzuTa Lin committed
419
420
        self.resize(size)
        self.move(position)
421
422
        saveDir = ustr(settings.get(SETTING_SAVE_DIR, None))
        self.lastOpenDir = ustr(settings.get(SETTING_LAST_OPEN_DIR, None))
tzutalin's avatar
tzutalin committed
423
        if saveDir is not None and os.path.exists(saveDir):
424
            self.defaultSaveDir = saveDir
425
426
            self.statusBar().showMessage('%s started. Annotation will be saved to %s' %
                                         (__appname__, self.defaultSaveDir))
427
428
            self.statusBar().show()

429
        self.restoreState(settings.get(SETTING_WIN_STATE, QByteArray()))
430
431
432
        Shape.line_color = self.lineColor = QColor(settings.get(SETTING_LINE_COLOR, DEFAULT_LINE_COLOR))
        Shape.fill_color = self.fillColor = QColor(settings.get(SETTING_FILL_COLOR, DEFAULT_FILL_COLOR))
        self.canvas.setDrawingColor(self.lineColor)
433
        # Add chris
tzutalin's avatar
tzutalin committed
434
        Shape.difficult = self.difficult
TzuTa Lin's avatar
TzuTa Lin committed
435

436
437
438
439
440
        def xbool(x):
            if isinstance(x, QVariant):
                return x.toBool()
            return bool(x)

441
        if xbool(settings.get(SETTING_ADVANCE_MODE, False)):
TzuTa Lin's avatar
TzuTa Lin committed
442
443
444
445
446
            self.actions.advancedMode.setChecked(True)
            self.toggleAdvancedMode()

        # Populate the File menu dynamically.
        self.updateFileMenu()
447
448
449
450
451
452

        # Since loading the file may take some time, make sure it runs in the background.
        if self.filePath and os.path.isdir(self.filePath):
            self.queueEvent(partial(self.importDirImages, self.filePath or ""))
        elif self.filePath:
            self.queueEvent(partial(self.loadFile, self.filePath or ""))
TzuTa Lin's avatar
TzuTa Lin committed
453
454
455
456
457
458

        # Callbacks:
        self.zoomWidget.valueChanged.connect(self.paintCanvas)

        self.populateModeActions()

459
460
461
462
        # Display cursor coordinates at the right of status bar
        self.labelCoordinates = QLabel('')
        self.statusBar().addPermanentWidget(self.labelCoordinates)

463
464
465
466
        # Open Dir if deafult file
        if self.filePath and os.path.isdir(self.filePath):
            self.openDirDialog(dirpath=self.filePath)

TzuTa Lin's avatar
TzuTa Lin committed
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
    ## Support Functions ##

    def noShapes(self):
        return not self.itemsToShapes

    def toggleAdvancedMode(self, value=True):
        self._beginner = not value
        self.canvas.setEditing(True)
        self.populateModeActions()
        self.editButton.setVisible(not value)
        if value:
            self.actions.createMode.setEnabled(True)
            self.actions.editMode.setEnabled(False)
            self.dock.setFeatures(self.dock.features() | self.dockFeatures)
        else:
            self.dock.setFeatures(self.dock.features() ^ self.dockFeatures)

    def populateModeActions(self):
        if self.beginner():
            tool, menu = self.actions.beginner, self.actions.beginnerContext
        else:
            tool, menu = self.actions.advanced, self.actions.advancedContext
        self.tools.clear()
        addActions(self.tools, tool)
        self.canvas.menus[0].clear()
        addActions(self.canvas.menus[0], menu)
        self.menus.edit.clear()
        actions = (self.actions.create,) if self.beginner()\
495
            else (self.actions.createMode, self.actions.editMode)
TzuTa Lin's avatar
TzuTa Lin committed
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
        addActions(self.menus.edit, actions + self.actions.editMenu)

    def setBeginner(self):
        self.tools.clear()
        addActions(self.tools, self.actions.beginner)

    def setAdvanced(self):
        self.tools.clear()
        addActions(self.tools, self.actions.advanced)

    def setDirty(self):
        self.dirty = True
        self.actions.save.setEnabled(True)

    def setClean(self):
        self.dirty = False
        self.actions.save.setEnabled(False)
        self.actions.create.setEnabled(True)

    def toggleActions(self, value=True):
        """Enable/Disable widgets which depend on an opened image."""
        for z in self.actions.zoomActions:
            z.setEnabled(value)
        for action in self.actions.onLoadActive:
            action.setEnabled(value)

    def queueEvent(self, function):
        QTimer.singleShot(0, function)

    def status(self, message, delay=5000):
        self.statusBar().showMessage(message, delay)

    def resetState(self):
        self.itemsToShapes.clear()
        self.shapesToItems.clear()
        self.labelList.clear()
532
        self.filePath = None
TzuTa Lin's avatar
TzuTa Lin committed
533
534
535
        self.imageData = None
        self.labelFile = None
        self.canvas.resetState()
536
        self.labelCoordinates.clear()
TzuTa Lin's avatar
TzuTa Lin committed
537
538
539
540
541
542
543

    def currentItem(self):
        items = self.labelList.selectedItems()
        if items:
            return items[0]
        return None

544
545
546
    def addRecentFile(self, filePath):
        if filePath in self.recentFiles:
            self.recentFiles.remove(filePath)
TzuTa Lin's avatar
TzuTa Lin committed
547
548
        elif len(self.recentFiles) >= self.maxRecent:
            self.recentFiles.pop()
549
        self.recentFiles.insert(0, filePath)
TzuTa Lin's avatar
TzuTa Lin committed
550
551
552
553
554
555
556
557

    def beginner(self):
        return self._beginner

    def advanced(self):
        return not self.beginner()

    ## Callbacks ##
tzutalin's avatar
tzutalin committed
558
    def showTutorialDialog(self):
TzuTa Lin's avatar
TzuTa Lin committed
559
560
        subprocess.Popen([self.screencastViewer, self.screencast])

tzutalin's avatar
tzutalin committed
561
562
563
564
    def showInfoDialog(self):
        msg = u'Name:{0} \nApp Version:{1} \n{2} '.format(__appname__, __version__, sys.version_info)
        QMessageBox.information(self, u'Information', msg)

TzuTa Lin's avatar
TzuTa Lin committed
565
566
567
568
569
570
571
572
573
574
    def createShape(self):
        assert self.beginner()
        self.canvas.setEditing(False)
        self.actions.create.setEnabled(False)

    def toggleDrawingSensitive(self, drawing=True):
        """In the middle of drawing, toggling between modes should be disabled."""
        self.actions.editMode.setEnabled(not drawing)
        if not drawing and self.beginner():
            # Cancel creation.
575
            print('Cancel creation.')
TzuTa Lin's avatar
TzuTa Lin committed
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
            self.canvas.setEditing(True)
            self.canvas.restoreCursor()
            self.actions.create.setEnabled(True)

    def toggleDrawMode(self, edit=True):
        self.canvas.setEditing(edit)
        self.actions.createMode.setEnabled(edit)
        self.actions.editMode.setEnabled(not edit)

    def setCreateMode(self):
        assert self.advanced()
        self.toggleDrawMode(False)

    def setEditMode(self):
        assert self.advanced()
        self.toggleDrawMode(True)
592
        self.labelSelectionChanged()
TzuTa Lin's avatar
TzuTa Lin committed
593
594

    def updateFileMenu(self):
595
        currFilePath = self.filePath
596

TzuTa Lin's avatar
TzuTa Lin committed
597
        def exists(filename):
598
            return os.path.exists(filename)
TzuTa Lin's avatar
TzuTa Lin committed
599
600
        menu = self.menus.recentFiles
        menu.clear()
601
602
        files = [f for f in self.recentFiles if f !=
                 currFilePath and exists(f)]
TzuTa Lin's avatar
TzuTa Lin committed
603
604
605
        for i, f in enumerate(files):
            icon = newIcon('labels')
            action = QAction(
606
                icon, '&%d %s' % (i + 1, QFileInfo(f).fileName()), self)
TzuTa Lin's avatar
TzuTa Lin committed
607
608
609
610
611
612
            action.triggered.connect(partial(self.loadRecent, f))
            menu.addAction(action)

    def popLabelListMenu(self, point):
        self.menus.labelList.exec_(self.labelList.mapToGlobal(point))

613
    def editLabel(self):
TzuTa Lin's avatar
TzuTa Lin committed
614
615
        if not self.canvas.editing():
            return
616
        item = self.currentItem()
TzuTa Lin's avatar
TzuTa Lin committed
617
618
619
        text = self.labelDialog.popUp(item.text())
        if text is not None:
            item.setText(text)
zhangjie's avatar
zhangjie committed
620
            item.setBackground(generateColorByText(text))
TzuTa Lin's avatar
TzuTa Lin committed
621
            self.setDirty()
622

623
624
    # Tzutalin 20160906 : Add file list and dock to move faster
    def fileitemDoubleClicked(self, item=None):
625
        currIndex = self.mImgList.index(ustr(item.text()))
626
        if currIndex < len(self.mImgList):
627
628
629
            filename = self.mImgList[currIndex]
            if filename:
                self.loadFile(filename)
630
631

    # Add chris
632
    def btnstate(self, item= None):
633
        """ Function to handle difficult examples
634
635
636
        Update on each object """
        if not self.canvas.editing():
            return
637

638
        item = self.currentItem()
639
        if not item: # If not selected Item, take the first one
640
641
642
643
644
645
646
647
            item = self.labelList.item(self.labelList.count()-1)

        difficult = self.diffcButton.isChecked()

        try:
            shape = self.itemsToShapes[item]
        except:
            pass
648
        # Checked and Update
649
650
651
652
653
654
655
656
        try:
            if difficult != shape.difficult:
                shape.difficult = difficult
                self.setDirty()
            else:  # User probably changed item visibility
                self.canvas.setShapeVisible(shape, item.checkState() == Qt.Checked)
        except:
            pass
TzuTa Lin's avatar
TzuTa Lin committed
657
658
659
660
661
662
663
664

    # React to canvas signals.
    def shapeSelectionChanged(self, selected=False):
        if self._noSelectionSlot:
            self._noSelectionSlot = False
        else:
            shape = self.canvas.selectedShape
            if shape:
665
                self.shapesToItems[shape].setSelected(True)
TzuTa Lin's avatar
TzuTa Lin committed
666
667
668
669
670
671
672
673
674
            else:
                self.labelList.clearSelection()
        self.actions.delete.setEnabled(selected)
        self.actions.copy.setEnabled(selected)
        self.actions.edit.setEnabled(selected)
        self.actions.shapeLineColor.setEnabled(selected)
        self.actions.shapeFillColor.setEnabled(selected)

    def addLabel(self, shape):
675
        item = HashableQListWidgetItem(shape.label)
TzuTa Lin's avatar
TzuTa Lin committed
676
677
        item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
        item.setCheckState(Qt.Checked)
678
        item.setBackground(generateColorByText(shape.label))
TzuTa Lin's avatar
TzuTa Lin committed
679
680
681
682
683
684
685
        self.itemsToShapes[item] = shape
        self.shapesToItems[shape] = item
        self.labelList.addItem(item)
        for action in self.actions.onShapesPresent:
            action.setEnabled(True)

    def remLabel(self, shape):
686
687
688
        if shape is None:
            # print('rm empty label')
            return
TzuTa Lin's avatar
TzuTa Lin committed
689
690
691
692
693
694
695
        item = self.shapesToItems[shape]
        self.labelList.takeItem(self.labelList.row(item))
        del self.shapesToItems[shape]
        del self.itemsToShapes[item]

    def loadLabels(self, shapes):
        s = []
696
        for label, points, line_color, fill_color, difficult in shapes:
TzuTa Lin's avatar
TzuTa Lin committed
697
698
699
            shape = Shape(label=label)
            for x, y in points:
                shape.addPoint(QPointF(x, y))
700
            shape.difficult = difficult
TzuTa Lin's avatar
TzuTa Lin committed
701
702
            shape.close()
            s.append(shape)
703

TzuTa Lin's avatar
TzuTa Lin committed
704
705
            if line_color:
                shape.line_color = QColor(*line_color)
706
707
708
            else:
                shape.line_color = generateColorByText(label)

TzuTa Lin's avatar
TzuTa Lin committed
709
710
            if fill_color:
                shape.fill_color = QColor(*fill_color)
711
712
            else:
                shape.fill_color = generateColorByText(label)
zhangjie's avatar
zhangjie committed
713

714
            self.addLabel(shape)
715

TzuTa Lin's avatar
TzuTa Lin committed
716
717
        self.canvas.loadShapes(s)

718
    def saveLabels(self, annotationFilePath):
719
        annotationFilePath = ustr(annotationFilePath)
Thibaut Mattio's avatar
Thibaut Mattio committed
720
721
722
        if self.labelFile is None:
            self.labelFile = LabelFile()
            self.labelFile.verified = self.canvas.verified
723

TzuTa Lin's avatar
TzuTa Lin committed
724
        def format_shape(s):
725
            return dict(label=s.label,
726
727
                        line_color=s.line_color.getRgb(),
                        fill_color=s.fill_color.getRgb(),
728
729
730
                        points=[(p.x(), p.y()) for p in s.points],
                       # add chris
                        difficult = s.difficult)
TzuTa Lin's avatar
TzuTa Lin committed
731
732

        shapes = [format_shape(shape) for shape in self.canvas.shapes]
733
        # Can add differrent annotation formats here
TzuTa Lin's avatar
TzuTa Lin committed
734
        try:
TzuTa Lin's avatar
Typo    
TzuTa Lin committed
735
            if self.usingPascalVocFormat is True:
736
                print ('Img: ' + self.filePath + ' -> Its xml: ' + annotationFilePath)
Thibaut Mattio's avatar
Thibaut Mattio committed
737
738
                self.labelFile.savePascalVocFormat(annotationFilePath, shapes, self.filePath, self.imageData,
                                                   self.lineColor.getRgb(), self.fillColor.getRgb())
TzuTa Lin's avatar
TzuTa Lin committed
739
            else:
Thibaut Mattio's avatar
Thibaut Mattio committed
740
741
                self.labelFile.save(annotationFilePath, shapes, self.filePath, self.imageData,
                                    self.lineColor.getRgb(), self.fillColor.getRgb())
TzuTa Lin's avatar
TzuTa Lin committed
742
            return True
743
        except LabelFileError as e:
744
            self.errorMessage(u'Error saving label data', u'<b>%s</b>' % e)
TzuTa Lin's avatar
TzuTa Lin committed
745
746
747
748
            return False

    def copySelectedShape(self):
        self.addLabel(self.canvas.copySelectedShape())
749
        # fix copy and delete
TzuTa Lin's avatar
TzuTa Lin committed
750
751
752
753
754
755
756
        self.shapeSelectionChanged(True)

    def labelSelectionChanged(self):
        item = self.currentItem()
        if item and self.canvas.editing():
            self._noSelectionSlot = True
            self.canvas.selectShape(self.itemsToShapes[item])
ChrisDal's avatar
ChrisDal committed
757
            shape = self.itemsToShapes[item]
758
759
            # Add Chris
            self.diffcButton.setChecked(shape.difficult)
TzuTa Lin's avatar
TzuTa Lin committed
760
761
762

    def labelItemChanged(self, item):
        shape = self.itemsToShapes[item]
763
        label = item.text()
TzuTa Lin's avatar
TzuTa Lin committed
764
        if label != shape.label:
765
            shape.label = item.text()
zhangjie's avatar
zhangjie committed
766
            shape.line_color = generateColorByText(shape.label)
TzuTa Lin's avatar
TzuTa Lin committed
767
            self.setDirty()
768
        else:  # User probably changed item visibility
TzuTa Lin's avatar
TzuTa Lin committed
769
770
            self.canvas.setShapeVisible(shape, item.checkState() == Qt.Checked)

771
    # Callback functions:
TzuTa Lin's avatar
TzuTa Lin committed
772
773
774
775
776
    def newShape(self):
        """Pop-up and give focus to the label editor.

        position MUST be in global coordinates.
        """
777
        if not self.useDefaultLabelCheckbox.isChecked() or not self.defaultLabelTextLine.text():
RegisWANG's avatar
RegisWANG committed
778
779
780
781
            if len(self.labelHist) > 0:
                self.labelDialog = LabelDialog(
                    parent=self, listItem=self.labelHist)

782
783
784
785
786
787
            # Sync single class mode from PR#106
            if self.singleClassMode.isChecked() and self.lastLabel:
                text = self.lastLabel
            else:
                text = self.labelDialog.popUp(text=self.prevLabelText)
                self.lastLabel = text
RegisWANG's avatar
RegisWANG committed
788
        else:
tzutalin's avatar
tzutalin committed
789
            text = self.defaultLabelTextLine.text()
790

791
792
        # Add Chris
        self.diffcButton.setChecked(False)
TzuTa Lin's avatar
TzuTa Lin committed
793
        if text is not None:
794
            self.prevLabelText = text
795
796
797
            generate_color = generateColorByText(text)
            shape = self.canvas.setLastLabel(text, generate_color, generate_color)
            self.addLabel(shape)
798
            if self.beginner():  # Switch to edit mode.
TzuTa Lin's avatar
TzuTa Lin committed
799
800
801
802
803
                self.canvas.setEditing(True)
                self.actions.create.setEnabled(True)
            else:
                self.actions.editMode.setEnabled(True)
            self.setDirty()
804
805
806

            if text not in self.labelHist:
                self.labelHist.append(text)
TzuTa Lin's avatar
TzuTa Lin committed
807
        else:
808
            # self.canvas.undoLastLine()
TzuTa Lin's avatar
TzuTa Lin committed
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
            self.canvas.resetAllLines()

    def scrollRequest(self, delta, orientation):
        units = - delta / (8 * 15)
        bar = self.scrollBars[orientation]
        bar.setValue(bar.value() + bar.singleStep() * units)

    def setZoom(self, value):
        self.actions.fitWidth.setChecked(False)
        self.actions.fitWindow.setChecked(False)
        self.zoomMode = self.MANUAL_ZOOM
        self.zoomWidget.setValue(value)

    def addZoom(self, increment=10):
        self.setZoom(self.zoomWidget.value() + increment)

    def zoomRequest(self, delta):
Lars Klein's avatar
Lars Klein committed
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
        # get the current scrollbar positions
        # calculate the percentages ~ coordinates
        h_bar = self.scrollBars[Qt.Horizontal]
        v_bar = self.scrollBars[Qt.Vertical]

        # get the current maximum, to know the difference after zooming
        h_bar_max = h_bar.maximum()
        v_bar_max = v_bar.maximum()

        # get the cursor position and canvas size
        # calculate the desired movement from 0 to 1
        # where 0 = move left
        #       1 = move right
        # up and down analogous
        cursor = QCursor()
        pos = cursor.pos()
842
        relative_pos = QWidget.mapFromGlobal(self, pos)
Lars Klein's avatar
Lars Klein committed
843

844
845
        cursor_x = relative_pos.x()
        cursor_y = relative_pos.y()
Lars Klein's avatar
Lars Klein committed
846
847
848
849

        w = self.scrollArea.width()
        h = self.scrollArea.height()

Lars Klein's avatar
Lars Klein committed
850
851
        # the scaling from 0 to 1 has some padding
        # you don't have to hit the very leftmost pixel for a maximum-left movement
Lars Klein's avatar
Lars Klein committed
852
853
854
855
        margin = 0.1
        move_x = (cursor_x - margin * w) / (w - 2 * margin * w)
        move_y = (cursor_y - margin * h) / (h - 2 * margin * h)

856
        # clamp the values from 0 to 1
Lars Klein's avatar
Lars Klein committed
857
858
859
        move_x = min(max(move_x, 0), 1)
        move_y = min(max(move_y, 0), 1)

Lars Klein's avatar
Lars Klein committed
860
        # zoom in
TzuTa Lin's avatar
TzuTa Lin committed
861
862
863
864
        units = delta / (8 * 15)
        scale = 10
        self.addZoom(scale * units)

Lars Klein's avatar
Lars Klein committed
865
866
        # get the difference in scrollbar values
        # this is how far we can move
Lars Klein's avatar
Lars Klein committed
867
868
869
870
871
872
873
874
875
876
        d_h_bar_max = h_bar.maximum() - h_bar_max
        d_v_bar_max = v_bar.maximum() - v_bar_max

        # get the new scrollbar values
        new_h_bar_value = h_bar.value() + move_x * d_h_bar_max
        new_v_bar_value = v_bar.value() + move_y * d_v_bar_max

        h_bar.setValue(new_h_bar_value)
        v_bar.setValue(new_v_bar_value)

TzuTa Lin's avatar
TzuTa Lin committed
877
878
879
880
881
882
883
884
885
886
887
888
889
    def setFitWindow(self, value=True):
        if value:
            self.actions.fitWidth.setChecked(False)
        self.zoomMode = self.FIT_WINDOW if value else self.MANUAL_ZOOM
        self.adjustScale()

    def setFitWidth(self, value=True):
        if value:
            self.actions.fitWindow.setChecked(False)
        self.zoomMode = self.FIT_WIDTH if value else self.MANUAL_ZOOM
        self.adjustScale()

    def togglePolygons(self, value):
890
        for item, shape in self.itemsToShapes.items():
TzuTa Lin's avatar
TzuTa Lin committed
891
892
            item.setCheckState(Qt.Checked if value else Qt.Unchecked)

893
    def loadFile(self, filePath=None):
TzuTa Lin's avatar
TzuTa Lin committed
894
895
896
        """Load the specified file, or the last opened file if None."""
        self.resetState()
        self.canvas.setEnabled(False)
897
        if filePath is None:
898
            filePath = self.settings.get(SETTING_FILENAME)
899

Tomas Raila's avatar
Tomas Raila committed
900
901
902
        # Make sure that filePath is a regular python string, rather than QString
        filePath = str(filePath)

903
        unicodeFilePath = ustr(filePath)
904
        # Tzutalin 20160906 : Add file list and dock to move faster
905
        # Highlight the file item
906
907
        if unicodeFilePath and self.fileListWidget.count() > 0:
            index = self.mImgList.index(unicodeFilePath)
908
            fileWidgetItem = self.fileListWidget.item(index)
909
            fileWidgetItem.setSelected(True)
910

911
912
        if unicodeFilePath and os.path.exists(unicodeFilePath):
            if LabelFile.isLabelFile(unicodeFilePath):
TzuTa Lin's avatar
TzuTa Lin committed
913
                try:
914
                    self.labelFile = LabelFile(unicodeFilePath)
915
                except LabelFileError as e:
TzuTa Lin's avatar
TzuTa Lin committed
916
                    self.errorMessage(u'Error opening file',
917
                                      (u"<p><b>%s</b></p>"
918
                                       u"<p>Make sure <i>%s</i> is a valid label file.")
919
920
                                      % (e, unicodeFilePath))
                    self.status("Error reading %s" % unicodeFilePath)
TzuTa Lin's avatar
TzuTa Lin committed
921
922
923
924
925
926
927
                    return False
                self.imageData = self.labelFile.imageData
                self.lineColor = QColor(*self.labelFile.lineColor)
                self.fillColor = QColor(*self.labelFile.fillColor)
            else:
                # Load image:
                # read data first and store for saving into label file.
928
                self.imageData = read(unicodeFilePath, None)
TzuTa Lin's avatar
TzuTa Lin committed
929
                self.labelFile = None
zhangjie's avatar
zhangjie committed
930

TzuTa Lin's avatar
TzuTa Lin committed
931
932
933
            image = QImage.fromData(self.imageData)
            if image.isNull():
                self.errorMessage(u'Error opening file',
934
935
                                  u"<p>Make sure <i>%s</i> is a valid image file." % unicodeFilePath)
                self.status("Error reading %s" % unicodeFilePath)
TzuTa Lin's avatar
TzuTa Lin committed
936
                return False
937
            self.status("Loaded %s" % os.path.basename(unicodeFilePath))
TzuTa Lin's avatar
TzuTa Lin committed
938
            self.image = image
939
            self.filePath = unicodeFilePath
TzuTa Lin's avatar
TzuTa Lin committed
940
941
942
943
944
945
946
            self.canvas.loadPixmap(QPixmap.fromImage(image))
            if self.labelFile:
                self.loadLabels(self.labelFile.shapes)
            self.setClean()
            self.canvas.setEnabled(True)
            self.adjustScale(initial=True)
            self.paintCanvas()
947
            self.addRecentFile(self.filePath)
TzuTa Lin's avatar
TzuTa Lin committed
948
            self.toggleActions(True)
949

950
            # Label xml file and show bound box according to its filename
951
952
953
954
955
956
957
            if self.usingPascalVocFormat is True:
                if self.defaultSaveDir is not None:
                    basename = os.path.basename(
                        os.path.splitext(self.filePath)[0]) + XML_EXT
                    xmlPath = os.path.join(self.defaultSaveDir, basename)
                    self.loadPascalXMLByFilename(xmlPath)
                else:
Jiye Qian's avatar
Jiye Qian committed
958
                    xmlPath = os.path.splitext(filePath)[0] + XML_EXT
959
960
                    if os.path.isfile(xmlPath):
                        self.loadPascalXMLByFilename(xmlPath)
961

962
            self.setWindowTitle(__appname__ + ' ' + filePath)
963
964

            # Default : select last item if there is at least one item
965
966
            if self.labelList.count():
                self.labelList.setCurrentItem(self.labelList.item(self.labelList.count()-1))
967
                self.labelList.item(self.labelList.count()-1).setSelected(True)
968

Thibaut Mattio's avatar
Thibaut Mattio committed
969
            self.canvas.setFocus(True)
TzuTa Lin's avatar
TzuTa Lin committed
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
            return True
        return False

    def resizeEvent(self, event):
        if self.canvas and not self.image.isNull()\
           and self.zoomMode != self.MANUAL_ZOOM:
            self.adjustScale()
        super(MainWindow, self).resizeEvent(event)

    def paintCanvas(self):
        assert not self.image.isNull(), "cannot paint null image"
        self.canvas.scale = 0.01 * self.zoomWidget.value()
        self.canvas.adjustSize()
        self.canvas.update()

    def adjustScale(self, initial=False):
        value = self.scalers[self.FIT_WINDOW if initial else self.zoomMode]()
        self.zoomWidget.setValue(int(100 * value))

    def scaleFitWindow(self):
        """Figure out the size of the pixmap in order to fit the main widget."""
991
        e = 2.0  # So that no scrollbars are generated.
TzuTa Lin's avatar
TzuTa Lin committed
992
993
        w1 = self.centralWidget().width() - e
        h1 = self.centralWidget().height() - e
994
        a1 = w1 / h1
TzuTa Lin's avatar
TzuTa Lin committed
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
        # Calculate a new scale value based on the pixmap's aspect ratio.
        w2 = self.canvas.pixmap.width() - 0.0
        h2 = self.canvas.pixmap.height() - 0.0
        a2 = w2 / h2
        return w1 / w2 if a2 >= a1 else h1 / h2

    def scaleFitWidth(self):
        # The epsilon does not seem to work too well here.
        w = self.centralWidget().width() - 2.0
        return w / self.canvas.pixmap.width()

    def closeEvent(self, event):
        if not self.mayContinue():
            event.ignore()
tzutalin's avatar
tzutalin committed
1009
        settings = self.settings
TzuTa Lin's avatar
TzuTa Lin committed
1010
1011
        # If it loads images from dir, don't load it at the begining
        if self.dirname is None:
tzutalin's avatar
tzutalin committed
1012
            settings[SETTING_FILENAME] = self.filePath if self.filePath else ''
TzuTa Lin's avatar
TzuTa Lin committed
1013
        else:
tzutalin's avatar
tzutalin committed
1014
1015
1016
1017
1018
1019
1020
1021
1022
            settings[SETTING_FILENAME] = ''

        settings[SETTING_WIN_SIZE] = self.size()
        settings[SETTING_WIN_POSE] = self.pos()
        settings[SETTING_WIN_STATE] = self.saveState()
        settings[SETTING_LINE_COLOR] = self.lineColor
        settings[SETTING_FILL_COLOR] = self.fillColor
        settings[SETTING_RECENT_FILES] = self.recentFiles
        settings[SETTING_ADVANCE_MODE] = not self._beginner
1023
        if self.defaultSaveDir and os.path.exists(self.defaultSaveDir):
tzutalin's avatar
tzutalin committed
1024
            settings[SETTING_SAVE_DIR] = ustr(self.defaultSaveDir)
TzuTa Lin's avatar
TzuTa Lin committed
1025
        else:
tzutalin's avatar
tzutalin committed
1026
            settings[SETTING_SAVE_DIR] = ""
1027

1028
        if self.lastOpenDir and os.path.exists(self.lastOpenDir):
tzutalin's avatar
tzutalin committed
1029
            settings[SETTING_LAST_OPEN_DIR] = self.lastOpenDir
1030
        else:
tzutalin's avatar
tzutalin committed
1031
            settings[SETTING_LAST_OPEN_DIR] = ""
1032

1033
1034
        settings[SETTING_AUTO_SAVE] = self.autoSaving.isChecked()
        settings[SETTING_SINGLE_CLASS] = self.singleClassMode.isChecked()
tzutalin's avatar
tzutalin committed
1035
        settings.save()
TzuTa Lin's avatar
TzuTa Lin committed
1036
1037
1038
1039
1040
1041
1042
    ## User Dialogs ##

    def loadRecent(self, filename):
        if self.mayContinue():
            self.loadFile(filename)

    def scanAllImages(self, folderPath):
1043
        extensions = ['.%s' % fmt.data().decode("ascii").lower() for fmt in QImageReader.supportedImageFormats()]
TzuTa Lin's avatar
TzuTa Lin committed
1044
1045
1046
1047
1048
        images = []

        for root, dirs, files in os.walk(folderPath):
            for file in files:
                if file.lower().endswith(tuple(extensions)):
1049
1050
                    relativePath = os.path.join(root, file)
                    path = ustr(os.path.abspath(relativePath))
1051
                    images.append(path)
1052
        images.sort(key=lambda x: x.lower())
TzuTa Lin's avatar
TzuTa Lin committed
1053
1054
        return images

1055
    def changeSavedirDialog(self, _value=False):
1056
        if self.defaultSaveDir is not None:
1057
            path = ustr(self.defaultSaveDir)
1058
        else:
1059
            path = '.'
TzuTa Lin's avatar
TzuTa Lin committed
1060

1061
        dirpath = ustr(QFileDialog.getExistingDirectory(self,
1062
                                                       '%s - Save annotations to the directory' % __appname__, path,  QFileDialog.ShowDirsOnly
1063
                                                       | QFileDialog.DontResolveSymlinks))
1064
1065
1066
1067

        if dirpath is not None and len(dirpath) > 1:
            self.defaultSaveDir = dirpath

1068
1069
        self.statusBar().showMessage('%s . Annotation will be saved to %s' %
                                     ('Change saved folder', self.defaultSaveDir))
1070
        self.statusBar().show()
TzuTa Lin's avatar
TzuTa Lin committed