labelImg.py 56.9 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
4
import distutils.spawn
TzuTa Lin's avatar
TzuTa Lin committed
5
import os.path
6
import platform
TzuTa Lin's avatar
TzuTa Lin committed
7
8
9
10
11
12
13
import re
import sys
import subprocess

from functools import partial
from collections import defaultdict

14
15
16
17
18
19
try:
    from PyQt5.QtGui import *
    from PyQt5.QtCore import *
    from PyQt5.QtWidgets import *
except ImportError:
    # needed for py3+qt4
20
21
    # Ref:
    # http://pyqt.sourceforge.net/Docs/PyQt4/incompatible_apis.html
22
    # http://stackoverflow.com/questions/21217399/pyqt4-qtcore-qvariant-object-instead-of-a-string
23
24
25
26
27
    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
28
29

import resources
tzutalin's avatar
tzutalin committed
30
# Add internal libs
31
from libs.constants import *
32
from libs.lib import struct, newAction, newIcon, addActions, fmtShortcut, generateColorByText
tzutalin's avatar
tzutalin committed
33
from libs.settings import Settings
34
35
36
37
38
39
40
41
42
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
Wang Yinghao's avatar
Wang Yinghao committed
43
from libs.yolo_io import YoloReader
Wang Yinghao's avatar
Wang Yinghao committed
44
from libs.yolo_io import TXT_EXT
45
from libs.ustr import ustr
tzutalin's avatar
tzutalin committed
46
from libs.version import __version__
TzuTa Lin's avatar
TzuTa Lin committed
47
48
49

__appname__ = 'labelImg'

50
51
# Utility functions and classes.

52
53
54
55
56
57
58
59
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
60
class WindowMixin(object):
61

TzuTa Lin's avatar
TzuTa Lin committed
62
63
64
65
66
67
68
69
70
    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)
71
        # toolbar.setOrientation(Qt.Vertical)
TzuTa Lin's avatar
TzuTa Lin committed
72
73
74
75
76
77
78
        toolbar.setToolButtonStyle(Qt.ToolButtonTextUnderIcon)
        if actions:
            addActions(toolbar, actions)
        self.addToolBar(Qt.LeftToolBarArea, toolbar)
        return toolbar


79
80
# PyQt5: TypeError: unhashable type: 'QListWidgetItem'
class HashableQListWidgetItem(QListWidgetItem):
81

82
83
    def __init__(self, *args):
        super(HashableQListWidgetItem, self).__init__(*args)
84

85
86
87
88
    def __hash__(self):
        return hash(id(self))


TzuTa Lin's avatar
TzuTa Lin committed
89
class MainWindow(QMainWindow, WindowMixin):
90
    FIT_WINDOW, FIT_WIDTH, MANUAL_ZOOM = list(range(3))
TzuTa Lin's avatar
TzuTa Lin committed
91

jeffrey's avatar
jeffrey committed
92
    def __init__(self, defaultFilename=None, defaultPrefdefClassFile=None, defaultSaveDir=None):
TzuTa Lin's avatar
TzuTa Lin committed
93
94
        super(MainWindow, self).__init__()
        self.setWindowTitle(__appname__)
95
96
97
98
99
100

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

TzuTa Lin's avatar
Typo    
TzuTa Lin committed
101
        # Save as Pascal voc xml
jeffrey's avatar
jeffrey committed
102
        self.defaultSaveDir = defaultSaveDir
TzuTa Lin's avatar
Typo    
TzuTa Lin committed
103
        self.usingPascalVocFormat = True
Wang Yinghao's avatar
Wang Yinghao committed
104
105
        self.usingYoloFormat = False

TzuTa Lin's avatar
TzuTa Lin committed
106
107
108
        # For loading all image under a directory
        self.mImgList = []
        self.dirname = None
109
        self.labelHist = []
110
        self.lastOpenDir = None
TzuTa Lin's avatar
TzuTa Lin committed
111
112
113
114
115
116

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

        self._noSelectionSlot = False
        self._beginner = True
117
        self.screencastViewer = self.getAvailableScreencastViewer()
TzuTa Lin's avatar
TzuTa Lin committed
118
        self.screencast = "https://youtu.be/p0nR2YsCY_U"
TzuTa Lin's avatar
TzuTa Lin committed
119

120
121
122
        # Load predefined classes to the list
        self.loadPredefinedClasses(defaultPrefdefClassFile)

TzuTa Lin's avatar
TzuTa Lin committed
123
        # Main widgets and related state.
124
        self.labelDialog = LabelDialog(parent=self, listItem=self.labelHist)
125

TzuTa Lin's avatar
TzuTa Lin committed
126
127
        self.itemsToShapes = {}
        self.shapesToItems = {}
128
        self.prevLabelText = ''
TzuTa Lin's avatar
TzuTa Lin committed
129
130
131

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

tzutalin's avatar
tzutalin committed
133
        # Create a widget for using default label
134
135
        self.useDefaultLabelCheckbox = QCheckBox(u'Use default label')
        self.useDefaultLabelCheckbox.setChecked(False)
tzutalin's avatar
tzutalin committed
136
        self.defaultLabelTextLine = QLineEdit()
137
138
139
140
141
        useDefaultLabelQHBoxLayout = QHBoxLayout()
        useDefaultLabelQHBoxLayout.addWidget(self.useDefaultLabelCheckbox)
        useDefaultLabelQHBoxLayout.addWidget(self.defaultLabelTextLine)
        useDefaultLabelContainer = QWidget()
        useDefaultLabelContainer.setLayout(useDefaultLabelQHBoxLayout)
tzutalin's avatar
tzutalin committed
142
143
144

        # Create a widget for edit and diffc button
        self.diffcButton = QCheckBox(u'difficult')
145
        self.diffcButton.setChecked(False)
146
        self.diffcButton.stateChanged.connect(self.btnstate)
tzutalin's avatar
tzutalin committed
147
148
        self.editButton = QToolButton()
        self.editButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
149

150
        # Add some of widgets to listLayout
tzutalin's avatar
tzutalin committed
151
        listLayout.addWidget(self.editButton)
152
        listLayout.addWidget(self.diffcButton)
153
        listLayout.addWidget(useDefaultLabelContainer)
154

tzutalin's avatar
tzutalin committed
155
156
157
158
159
160
161
162
163
164
        # 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
165
166
167

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

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

TzuTa Lin's avatar
TzuTa Lin committed
182
183
184
        self.zoomWidget = ZoomWidget()
        self.colorDialog = ColorDialog(parent=self)

185
        self.canvas = Canvas(parent=self)
TzuTa Lin's avatar
TzuTa Lin committed
186
187
188
189
190
191
192
193
        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()
194
        }
Lars Klein's avatar
Lars Klein committed
195
        self.scrollArea = scroll
TzuTa Lin's avatar
TzuTa Lin committed
196
197
198
199
200
201
202
203
204
        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)
205
206
        # Tzutalin 20160906 : Add file list and dock to move faster
        self.addDockWidget(Qt.RightDockWidgetArea, self.filedock)
tzutalin's avatar
tzutalin committed
207
        self.filedock.setFeatures(QDockWidget.DockWidgetFloatable)
208

tzutalin's avatar
tzutalin committed
209
        self.dockFeatures = QDockWidget.DockWidgetClosable | QDockWidget.DockWidgetFloatable
TzuTa Lin's avatar
TzuTa Lin committed
210
211
212
213
214
        self.dock.setFeatures(self.dock.features() ^ self.dockFeatures)

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

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

220
        opendir = action('&Open Dir', self.openDirDialog,
221
                         'Ctrl+u', 'open', u'Open Dir')
TzuTa Lin's avatar
TzuTa Lin committed
222

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

226
        openAnnotation = action('&Open Annotation', self.openAnnotationDialog,
227
                                'Ctrl+Shift+O', 'open', u'Open Annotation')
228

TzuTa Lin's avatar
TzuTa Lin committed
229
        openNextImg = action('&Next Image', self.openNextImg,
230
                             'd', 'next', u'Open Next')
TzuTa Lin's avatar
TzuTa Lin committed
231

tzutalin's avatar
tzutalin committed
232
        openPrevImg = action('&Prev Image', self.openPrevImg,
233
                             'a', 'prev', u'Open Prev')
tzutalin's avatar
tzutalin committed
234

Thibaut Mattio's avatar
Thibaut Mattio committed
235
236
237
        verify = action('&Verify Image', self.verifyImg,
                        'space', 'verify', u'Verify Image')

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

Wang Yinghao's avatar
Wang Yinghao committed
241
        save_format = action('&PascalVOC', self.change_format,
Wang Yinghao's avatar
Wang Yinghao committed
242
                      'Ctrl+', 'format_voc', u'Change save format', enabled=True)
Wang Yinghao's avatar
Wang Yinghao committed
243

TzuTa Lin's avatar
TzuTa Lin committed
244
        saveAs = action('&Save As', self.saveFileAs,
tzutalin's avatar
tzutalin committed
245
246
247
                        '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
248

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

251
        color1 = action('Box Line Color', self.chooseColor1,
252
                        'Ctrl+L', 'color_line', u'Choose Box line color')
TzuTa Lin's avatar
TzuTa Lin committed
253
254

        createMode = action('Create\nRectBox', self.setCreateMode,
255
                            'w', 'new', u'Start drawing Boxs', enabled=False)
TzuTa Lin's avatar
TzuTa Lin committed
256
        editMode = action('&Edit\nRectBox', self.setEditMode,
257
                          'Ctrl+J', 'edit', u'Move and edit Boxs', enabled=False)
TzuTa Lin's avatar
TzuTa Lin committed
258
259

        create = action('Create\nRectBox', self.createShape,
260
                        'w', 'new', u'Draw a new Box', enabled=False)
TzuTa Lin's avatar
TzuTa Lin committed
261
        delete = action('Delete\nRectBox', self.deleteSelectedShape,
262
                        'Delete', 'delete', u'Delete', enabled=False)
TzuTa Lin's avatar
TzuTa Lin committed
263
        copy = action('&Duplicate\nRectBox', self.copySelectedShape,
264
265
                      'Ctrl+D', 'copy', u'Create a duplicate of the selected Box',
                      enabled=False)
TzuTa Lin's avatar
TzuTa Lin committed
266
267

        advancedMode = action('&Advanced Mode', self.toggleAdvancedMode,
268
269
                              'Ctrl+Shift+A', 'expert', u'Switch to advanced mode',
                              checkable=True)
TzuTa Lin's avatar
TzuTa Lin committed
270
271

        hideAll = action('&Hide\nRectBox', partial(self.togglePolygons, False),
272
273
                         'Ctrl+H', 'hide', u'Hide all Boxs',
                         enabled=False)
TzuTa Lin's avatar
TzuTa Lin committed
274
        showAll = action('&Show\nRectBox', partial(self.togglePolygons, True),
275
276
                         'Ctrl+A', 'hide', u'Show all Boxs',
                         enabled=False)
TzuTa Lin's avatar
TzuTa Lin committed
277

tzutalin's avatar
tzutalin committed
278
279
        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
280
281
282
283

        zoom = QWidgetAction(self)
        zoom.setDefaultWidget(self.zoomWidget)
        self.zoomWidget.setWhatsThis(
284
285
286
            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
287
288
289
        self.zoomWidget.setEnabled(False)

        zoomIn = action('Zoom &In', partial(self.addZoom, 10),
290
                        'Ctrl++', 'zoom-in', u'Increase zoom level', enabled=False)
TzuTa Lin's avatar
TzuTa Lin committed
291
        zoomOut = action('&Zoom Out', partial(self.addZoom, -10),
292
                         'Ctrl+-', 'zoom-out', u'Decrease zoom level', enabled=False)
TzuTa Lin's avatar
TzuTa Lin committed
293
        zoomOrg = action('&Original size', partial(self.setZoom, 100),
294
                         'Ctrl+=', 'zoom', u'Zoom to original size', enabled=False)
TzuTa Lin's avatar
TzuTa Lin committed
295
        fitWindow = action('&Fit Window', self.setFitWindow,
296
297
                           'Ctrl+F', 'fit-window', u'Zoom follows window size',
                           checkable=True, enabled=False)
TzuTa Lin's avatar
TzuTa Lin committed
298
        fitWidth = action('Fit &Width', self.setFitWidth,
299
                          'Ctrl+Shift+F', 'fit-width', u'Zoom follows window width',
300
                          checkable=True, enabled=False)                      
TzuTa Lin's avatar
TzuTa Lin committed
301
        # Group zoom controls into a list for easier toggling.
302
303
        zoomActions = (self.zoomWidget, zoomIn, zoomOut,
                       zoomOrg, fitWindow, fitWidth)
TzuTa Lin's avatar
TzuTa Lin committed
304
305
306
307
308
309
310
311
312
        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,
313
314
                      'Ctrl+E', 'edit', u'Modify the label of the selected Box',
                      enabled=False)
TzuTa Lin's avatar
TzuTa Lin committed
315
316
317
        self.editButton.setDefaultAction(edit)

        shapeLineColor = action('Shape &Line Color', self.chshapeLineColor,
318
319
                                icon='color_line', tip=u'Change the line color for this specific shape',
                                enabled=False)
TzuTa Lin's avatar
TzuTa Lin committed
320
        shapeFillColor = action('Shape &Fill Color', self.chshapeFillColor,
321
322
                                icon='color', tip=u'Change the fill color for this specific shape',
                                enabled=False)
TzuTa Lin's avatar
TzuTa Lin committed
323
324
325
326
327
328
329
330
331

        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)
332
333
        self.labelList.customContextMenuRequested.connect(
            self.popLabelListMenu)
TzuTa Lin's avatar
TzuTa Lin committed
334
335

        # Store actions for further handling.
Wang Yinghao's avatar
Wang Yinghao committed
336
        self.actions = struct(save=save, save_format=save_format, saveAs=saveAs, open=open, close=close, resetAll = resetAll,
337
                              lineColor=color1, create=create, delete=delete, edit=edit, copy=copy,
338
339
340
341
342
343
                              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
344
                                  open, opendir, save, saveAs, close, resetAll, quit),
345
346
                              beginner=(), advanced=(),
                              editMenu=(edit, copy, delete,
347
                                        None, color1),
348
349
350
351
352
353
                              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
354
355

        self.menus = struct(
356
357
358
359
360
361
            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
362

363
        # Auto saving : Enable auto saving if pressing next
364
365
        self.autoSaving = QAction("Auto Saving", self)
        self.autoSaving.setCheckable(True)
366
        self.autoSaving.setChecked(settings.get(SETTING_AUTO_SAVE, False))
367
368
369
370
        # Sync single class mode from PR#106
        self.singleClassMode = QAction("Single Class Mode", self)
        self.singleClassMode.setShortcut("Ctrl+Shift+S")
        self.singleClassMode.setCheckable(True)
371
        self.singleClassMode.setChecked(settings.get(SETTING_SINGLE_CLASS, False))
372
        self.lastLabel = None
373
374
375
376
        # Add option to enable/disable labels being painted at the top of bounding boxes
        self.paintLabelsOption = QAction("Paint Labels", self)
        self.paintLabelsOption.setShortcut("Ctrl+Shift+P")
        self.paintLabelsOption.setCheckable(True)
tzutalin's avatar
tzutalin committed
377
        self.paintLabelsOption.setChecked(settings.get(SETTING_PAINT_LABEL, False))
378
        self.paintLabelsOption.triggered.connect(self.togglePaintLabelsOption)
379

TzuTa Lin's avatar
TzuTa Lin committed
380
        addActions(self.menus.file,
Wang Yinghao's avatar
Wang Yinghao committed
381
                   (open, opendir, changeSavedir, openAnnotation, self.menus.recentFiles, save, save_format, saveAs, close, resetAll, quit))
tzutalin's avatar
tzutalin committed
382
        addActions(self.menus.help, (help, showInfo))
TzuTa Lin's avatar
TzuTa Lin committed
383
        addActions(self.menus.view, (
384
385
            self.autoSaving,
            self.singleClassMode,
386
            self.paintLabelsOption,
TzuTa Lin's avatar
TzuTa Lin committed
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
            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 = (
Wang Yinghao's avatar
Wang Yinghao committed
402
            open, opendir, changeSavedir, openNextImg, openPrevImg, verify, save, save_format, None, create, copy, delete, None,
TzuTa Lin's avatar
TzuTa Lin committed
403
404
405
            zoomIn, zoom, zoomOut, fitWindow, fitWidth)

        self.actions.advanced = (
Wang Yinghao's avatar
Wang Yinghao committed
406
            open, opendir, changeSavedir, openNextImg, openPrevImg, save, save_format, None,
TzuTa Lin's avatar
TzuTa Lin committed
407
408
409
410
411
412
413
414
            createMode, editMode, None,
            hideAll, showAll)

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

        # Application state.
        self.image = QImage()
415
        self.filePath = ustr(defaultFilename)
TzuTa Lin's avatar
TzuTa Lin committed
416
417
418
419
420
421
        self.recentFiles = []
        self.maxRecent = 7
        self.lineColor = None
        self.fillColor = None
        self.zoom_level = 100
        self.fit_window = False
422
        # Add Chris
423
        self.difficult = False
TzuTa Lin's avatar
TzuTa Lin committed
424

425
426
427
428
429
430
431
432
433
434
        ## 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
435
436
        self.resize(size)
        self.move(position)
437
438
        saveDir = ustr(settings.get(SETTING_SAVE_DIR, None))
        self.lastOpenDir = ustr(settings.get(SETTING_LAST_OPEN_DIR, None))
jeffrey's avatar
jeffrey committed
439
        if self.defaultSaveDir is None and saveDir is not None and os.path.exists(saveDir):
440
            self.defaultSaveDir = saveDir
441
442
            self.statusBar().showMessage('%s started. Annotation will be saved to %s' %
                                         (__appname__, self.defaultSaveDir))
443
444
            self.statusBar().show()

445
        self.restoreState(settings.get(SETTING_WIN_STATE, QByteArray()))
446
447
448
        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)
449
        # Add chris
tzutalin's avatar
tzutalin committed
450
        Shape.difficult = self.difficult
TzuTa Lin's avatar
TzuTa Lin committed
451

452
453
454
455
456
        def xbool(x):
            if isinstance(x, QVariant):
                return x.toBool()
            return bool(x)

457
        if xbool(settings.get(SETTING_ADVANCE_MODE, False)):
TzuTa Lin's avatar
TzuTa Lin committed
458
459
460
461
462
            self.actions.advancedMode.setChecked(True)
            self.toggleAdvancedMode()

        # Populate the File menu dynamically.
        self.updateFileMenu()
463
464
465
466
467
468

        # 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
469
470
471
472
473
474

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

        self.populateModeActions()

475
476
477
478
        # Display cursor coordinates at the right of status bar
        self.labelCoordinates = QLabel('')
        self.statusBar().addPermanentWidget(self.labelCoordinates)

479
480
481
482
        # 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
483
    ## Support Functions ##
Wang Yinghao's avatar
Wang Yinghao committed
484
    def set_format(self, save_format):
tzutalin's avatar
tzutalin committed
485
486
        if save_format == FORMAT_PASCALVOC:
            self.actions.save_format.setText(FORMAT_PASCALVOC)
Wang Yinghao's avatar
Wang Yinghao committed
487
            self.actions.save_format.setIcon(newIcon("format_voc"))
Wang Yinghao's avatar
Wang Yinghao committed
488
489
            self.usingPascalVocFormat = True
            self.usingYoloFormat = False
vdalv's avatar
vdalv committed
490
            LabelFile.suffix = XML_EXT
Wang Yinghao's avatar
Wang Yinghao committed
491

tzutalin's avatar
tzutalin committed
492
493
        elif save_format == FORMAT_YOLO:
            self.actions.save_format.setText(FORMAT_YOLO)
Wang Yinghao's avatar
Wang Yinghao committed
494
            self.actions.save_format.setIcon(newIcon("format_yolo"))
Wang Yinghao's avatar
Wang Yinghao committed
495
496
            self.usingPascalVocFormat = False
            self.usingYoloFormat = True
vdalv's avatar
vdalv committed
497
            LabelFile.suffix = TXT_EXT
TzuTa Lin's avatar
TzuTa Lin committed
498

Wang Yinghao's avatar
Wang Yinghao committed
499
    def change_format(self):
tzutalin's avatar
tzutalin committed
500
501
        if self.usingPascalVocFormat: self.set_format(FORMAT_YOLO)
        elif self.usingYoloFormat: self.set_format(FORMAT_PASCALVOC)
Wang Yinghao's avatar
Wang Yinghao committed
502

TzuTa Lin's avatar
TzuTa Lin committed
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
    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()\
529
            else (self.actions.createMode, self.actions.editMode)
TzuTa Lin's avatar
TzuTa Lin committed
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
        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()
566
        self.filePath = None
TzuTa Lin's avatar
TzuTa Lin committed
567
568
569
        self.imageData = None
        self.labelFile = None
        self.canvas.resetState()
570
        self.labelCoordinates.clear()
TzuTa Lin's avatar
TzuTa Lin committed
571
572
573
574
575
576
577

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

578
579
580
    def addRecentFile(self, filePath):
        if filePath in self.recentFiles:
            self.recentFiles.remove(filePath)
TzuTa Lin's avatar
TzuTa Lin committed
581
582
        elif len(self.recentFiles) >= self.maxRecent:
            self.recentFiles.pop()
583
        self.recentFiles.insert(0, filePath)
TzuTa Lin's avatar
TzuTa Lin committed
584
585
586
587
588
589
590

    def beginner(self):
        return self._beginner

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

591
592
593
594
595
596
597
598
599
600
    def getAvailableScreencastViewer(self):
        osName = platform.system()

        if osName == 'Windows':
            return ['C:\\Program Files\\Internet Explorer\\iexplore.exe']
        elif osName == 'Linux':
            return ['xdg-open']
        elif osName == 'Darwin':
            return ['open', '-a', 'Safari']

TzuTa Lin's avatar
TzuTa Lin committed
601
    ## Callbacks ##
tzutalin's avatar
tzutalin committed
602
    def showTutorialDialog(self):
603
        subprocess.Popen(self.screencastViewer + [self.screencast])
TzuTa Lin's avatar
TzuTa Lin committed
604

tzutalin's avatar
tzutalin committed
605
606
607
608
    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
609
610
611
612
613
614
615
616
617
618
    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.
619
            print('Cancel creation.')
TzuTa Lin's avatar
TzuTa Lin committed
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
            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)
636
        self.labelSelectionChanged()
TzuTa Lin's avatar
TzuTa Lin committed
637
638

    def updateFileMenu(self):
639
        currFilePath = self.filePath
640

TzuTa Lin's avatar
TzuTa Lin committed
641
        def exists(filename):
642
            return os.path.exists(filename)
TzuTa Lin's avatar
TzuTa Lin committed
643
644
        menu = self.menus.recentFiles
        menu.clear()
645
646
        files = [f for f in self.recentFiles if f !=
                 currFilePath and exists(f)]
TzuTa Lin's avatar
TzuTa Lin committed
647
648
649
        for i, f in enumerate(files):
            icon = newIcon('labels')
            action = QAction(
650
                icon, '&%d %s' % (i + 1, QFileInfo(f).fileName()), self)
TzuTa Lin's avatar
TzuTa Lin committed
651
652
653
654
655
656
            action.triggered.connect(partial(self.loadRecent, f))
            menu.addAction(action)

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

657
    def editLabel(self):
TzuTa Lin's avatar
TzuTa Lin committed
658
659
        if not self.canvas.editing():
            return
660
        item = self.currentItem()
TzuTa Lin's avatar
TzuTa Lin committed
661
662
663
        text = self.labelDialog.popUp(item.text())
        if text is not None:
            item.setText(text)
zhangjie's avatar
zhangjie committed
664
            item.setBackground(generateColorByText(text))
TzuTa Lin's avatar
TzuTa Lin committed
665
            self.setDirty()
666

667
668
    # Tzutalin 20160906 : Add file list and dock to move faster
    def fileitemDoubleClicked(self, item=None):
669
        currIndex = self.mImgList.index(ustr(item.text()))
670
        if currIndex < len(self.mImgList):
671
672
673
            filename = self.mImgList[currIndex]
            if filename:
                self.loadFile(filename)
674
675

    # Add chris
676
    def btnstate(self, item= None):
677
        """ Function to handle difficult examples
678
679
680
        Update on each object """
        if not self.canvas.editing():
            return
681

682
        item = self.currentItem()
683
        if not item: # If not selected Item, take the first one
684
685
686
687
688
689
690
691
            item = self.labelList.item(self.labelList.count()-1)

        difficult = self.diffcButton.isChecked()

        try:
            shape = self.itemsToShapes[item]
        except:
            pass
692
        # Checked and Update
693
694
695
696
697
698
699
700
        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
701
702
703
704
705
706
707
708

    # React to canvas signals.
    def shapeSelectionChanged(self, selected=False):
        if self._noSelectionSlot:
            self._noSelectionSlot = False
        else:
            shape = self.canvas.selectedShape
            if shape:
709
                self.shapesToItems[shape].setSelected(True)
TzuTa Lin's avatar
TzuTa Lin committed
710
711
712
713
714
715
716
717
718
            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):
719
        shape.paintLabel = self.paintLabelsOption.isChecked()
720
        item = HashableQListWidgetItem(shape.label)
TzuTa Lin's avatar
TzuTa Lin committed
721
722
        item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
        item.setCheckState(Qt.Checked)
723
        item.setBackground(generateColorByText(shape.label))
TzuTa Lin's avatar
TzuTa Lin committed
724
725
726
727
728
729
730
        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):
731
732
733
        if shape is None:
            # print('rm empty label')
            return
TzuTa Lin's avatar
TzuTa Lin committed
734
735
736
737
738
739
740
        item = self.shapesToItems[shape]
        self.labelList.takeItem(self.labelList.row(item))
        del self.shapesToItems[shape]
        del self.itemsToShapes[item]

    def loadLabels(self, shapes):
        s = []
741
        for label, points, box, line_color, fill_color, difficult in shapes:
TzuTa Lin's avatar
TzuTa Lin committed
742
            shape = Shape(label=label)
743
744
745
746
747
748
749
750
751
            center = QPointF(*box[0])
            width, height = box[1]
            angle = box[2]
            shape.addPoint(QPointF(center.x() - width / 2, center.y() - height / 2))
            shape.addPoint(QPointF(center.x() + width / 2, center.y() - height / 2))
            shape.addPoint(QPointF(center.x() + width / 2, center.y() + height / 2))
            shape.addPoint(QPointF(center.x() - width / 2, center.y() + height / 2))
            shape.applyRotationAngle(angle, center)
            shape.updateState()
752
            shape.difficult = difficult
TzuTa Lin's avatar
TzuTa Lin committed
753
754
            shape.close()
            s.append(shape)
755
756
            shape.line_color = QColor(*line_color) if line_color else generateColorByText(label)
            shape.fill_color = QColor(*fill_color) if fill_color else generateColorByText(label)
757
            self.addLabel(shape)
758

TzuTa Lin's avatar
TzuTa Lin committed
759
760
        self.canvas.loadShapes(s)

761
    def saveLabels(self, annotationFilePath):
762
        annotationFilePath = ustr(annotationFilePath)
Thibaut Mattio's avatar
Thibaut Mattio committed
763
764
765
        if self.labelFile is None:
            self.labelFile = LabelFile()
            self.labelFile.verified = self.canvas.verified
766

TzuTa Lin's avatar
TzuTa Lin committed
767
        def format_shape(s):
768
            s.updateState()
769
            return dict(label=s.label,
770
771
                        line_color=s.line_color.getRgb(),
                        fill_color=s.fill_color.getRgb(),
772
773
                        points=[(p.x(), p.y()) for p in s.points],
                       # add chris
774
775
776
777
778
                        difficult=s.difficult,
                        center=[s.center.x(), s.center.y()],
                        width=s.width,
                        height=s.height,
                        angle=s.currentAngle)
TzuTa Lin's avatar
TzuTa Lin committed
779
780

        shapes = [format_shape(shape) for shape in self.canvas.shapes]
781
        # Can add differrent annotation formats here
TzuTa Lin's avatar
TzuTa Lin committed
782
        try:
TzuTa Lin's avatar
Typo    
TzuTa Lin committed
783
            if self.usingPascalVocFormat is True:
tzutalin's avatar
tzutalin committed
784
                if ustr(annotationFilePath[-4:]) != ".xml":
785
                    annotationFilePath += XML_EXT
786
                print ('Img: ' + self.filePath + ' -> Its xml: ' + annotationFilePath)
Thibaut Mattio's avatar
Thibaut Mattio committed
787
788
                self.labelFile.savePascalVocFormat(annotationFilePath, shapes, self.filePath, self.imageData,
                                                   self.lineColor.getRgb(), self.fillColor.getRgb())
Wang Yinghao's avatar
Wang Yinghao committed
789
            elif self.usingYoloFormat is True:
790
791
                if annotationFilePath[-4:] != ".txt":
                    annotationFilePath += TXT_EXT
Wang Yinghao's avatar
Wang Yinghao committed
792
793
794
                print ('Img: ' + self.filePath + ' -> Its txt: ' + annotationFilePath)
                self.labelFile.saveYoloFormat(annotationFilePath, shapes, self.filePath, self.imageData, self.labelHist,
                                                   self.lineColor.getRgb(), self.fillColor.getRgb())
TzuTa Lin's avatar
TzuTa Lin committed
795
            else:
Thibaut Mattio's avatar
Thibaut Mattio committed
796
797
                self.labelFile.save(annotationFilePath, shapes, self.filePath, self.imageData,
                                    self.lineColor.getRgb(), self.fillColor.getRgb())
TzuTa Lin's avatar
TzuTa Lin committed
798
            return True
799
        except LabelFileError as e:
800
            self.errorMessage(u'Error saving label data', u'<b>%s</b>' % e)
TzuTa Lin's avatar
TzuTa Lin committed
801
802
803
804
            return False

    def copySelectedShape(self):
        self.addLabel(self.canvas.copySelectedShape())
805
        # fix copy and delete
TzuTa Lin's avatar
TzuTa Lin committed
806
807
808
809
810
811
812
        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
813
            shape = self.itemsToShapes[item]
814
815
            # Add Chris
            self.diffcButton.setChecked(shape.difficult)
TzuTa Lin's avatar
TzuTa Lin committed
816
817
818

    def labelItemChanged(self, item):
        shape = self.itemsToShapes[item]
819
        label = item.text()
TzuTa Lin's avatar
TzuTa Lin committed
820
        if label != shape.label:
821
            shape.label = item.text()
zhangjie's avatar
zhangjie committed
822
            shape.line_color = generateColorByText(shape.label)
TzuTa Lin's avatar
TzuTa Lin committed
823
            self.setDirty()
824
        else:  # User probably changed item visibility
TzuTa Lin's avatar
TzuTa Lin committed
825
826
            self.canvas.setShapeVisible(shape, item.checkState() == Qt.Checked)

827
    # Callback functions:
TzuTa Lin's avatar
TzuTa Lin committed
828
829
830
831
832
    def newShape(self):
        """Pop-up and give focus to the label editor.

        position MUST be in global coordinates.
        """
833
        if not self.useDefaultLabelCheckbox.isChecked() or not self.defaultLabelTextLine.text():
RegisWANG's avatar
RegisWANG committed
834
835
836
837
            if len(self.labelHist) > 0:
                self.labelDialog = LabelDialog(
                    parent=self, listItem=self.labelHist)

838
839
840
841
842
843
            # 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
844
        else:
tzutalin's avatar
tzutalin committed
845
            text = self.defaultLabelTextLine.text()
846

847
848
        # Add Chris
        self.diffcButton.setChecked(False)
TzuTa Lin's avatar
TzuTa Lin committed
849
        if text is not None:
850
            self.prevLabelText = text
851
852
853
            generate_color = generateColorByText(text)
            shape = self.canvas.setLastLabel(text, generate_color, generate_color)
            self.addLabel(shape)
854
            if self.beginner():  # Switch to edit mode.
TzuTa Lin's avatar
TzuTa Lin committed
855
856
857
858
859
                self.canvas.setEditing(True)
                self.actions.create.setEnabled(True)
            else:
                self.actions.editMode.setEnabled(True)
            self.setDirty()
860
861
862

            if text not in self.labelHist:
                self.labelHist.append(text)
TzuTa Lin's avatar
TzuTa Lin committed
863
        else:
864
            # self.canvas.undoLastLine()
TzuTa Lin's avatar
TzuTa Lin committed
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
            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
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
        # 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()
898
        relative_pos = QWidget.mapFromGlobal(self, pos)
Lars Klein's avatar
Lars Klein committed
899

900
901
        cursor_x = relative_pos.x()
        cursor_y = relative_pos.y()
Lars Klein's avatar
Lars Klein committed
902
903
904
905

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

Lars Klein's avatar
Lars Klein committed
906
907
        # 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
908
909
910
911
        margin = 0.1
        move_x = (cursor_x - margin * w) / (w - 2 * margin * w)
        move_y = (cursor_y - margin * h) / (h - 2 * margin * h)

912
        # clamp the values from 0 to 1
Lars Klein's avatar
Lars Klein committed
913
914
915
        move_x = min(max(move_x, 0), 1)
        move_y = min(max(move_y, 0), 1)

Lars Klein's avatar
Lars Klein committed
916
        # zoom in
TzuTa Lin's avatar
TzuTa Lin committed
917
918
919
920
        units = delta / (8 * 15)
        scale = 10
        self.addZoom(scale * units)

Lars Klein's avatar
Lars Klein committed
921
922
        # get the difference in scrollbar values
        # this is how far we can move
Lars Klein's avatar
Lars Klein committed
923
924
925
926
927
928
929
930
931
932
        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
933
934
935
936
937
938
939
940
941
942
943
944
945
    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):
946
        for item, shape in self.itemsToShapes.items():
TzuTa Lin's avatar
TzuTa Lin committed
947
948
            item.setCheckState(Qt.Checked if value else Qt.Unchecked)

949
    def loadFile(self, filePath=None):
TzuTa Lin's avatar
TzuTa Lin committed
950
951
952
        """Load the specified file, or the last opened file if None."""
        self.resetState()
        self.canvas.setEnabled(False)
953
        if filePath is None:
954
            filePath = self.settings.get(SETTING_FILENAME)
955

Tomas Raila's avatar
Tomas Raila committed
956
        # Make sure that filePath is a regular python string, rather than QString
tzutalin's avatar
tzutalin committed
957
        filePath = ustr(filePath)
Tomas Raila's avatar
Tomas Raila committed
958

959
        unicodeFilePath = ustr(filePath)
960
        # Tzutalin 20160906 : Add file list and dock to move faster
961
        # Highlight the file item
962
963
        if unicodeFilePath and self.fileListWidget.count() > 0:
            index = self.mImgList.index(unicodeFilePath)
964
            fileWidgetItem = self.fileListWidget.item(index)
965
            fileWidgetItem.setSelected(True)
966

967
968
        if unicodeFilePath and os.path.exists(unicodeFilePath):
            if LabelFile.isLabelFile(unicodeFilePath):
TzuTa Lin's avatar
TzuTa Lin committed
969
                try:
970
                    self.labelFile = LabelFile(unicodeFilePath)
971
                except LabelFileError as e:
TzuTa Lin's avatar
TzuTa Lin committed
972
                    self.errorMessage(u'Error opening file',
973
                                      (u"<p><b>%s</b></p>"
974
                                       u"<p>Make sure <i>%s</i> is a valid label file.")
975
976
                                      % (e, unicodeFilePath))
                    self.status("Error reading %s" % unicodeFilePath)
TzuTa Lin's avatar
TzuTa Lin committed
977
978
979
980
                    return False
                self.imageData = self.labelFile.imageData
                self.lineColor = QColor(*self.labelFile.lineColor)
                self.fillColor = QColor(*self.labelFile.fillColor)
vdalv's avatar
vdalv committed
981
                self.canvas.verified = self.labelFile.verified
TzuTa Lin's avatar
TzuTa Lin committed
982
983
984
            else:
                # Load image:
                # read data first and store for saving into label file.
985
                self.imageData = read(unicodeFilePath, None)
TzuTa Lin's avatar
TzuTa Lin committed
986
                self.labelFile = None
vdalv's avatar
vdalv committed
987
                self.canvas.verified = False
zhangjie's avatar
zhangjie committed
988

TzuTa Lin's avatar
TzuTa Lin committed
989
990
991
            image = QImage.fromData(self.imageData)
            if image.isNull():
                self.errorMessage(u'Error opening file',
992
993
                                  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
994
                return False
995
            self.status("Loaded %s" % os.path.basename(unicodeFilePath))
TzuTa Lin's avatar
TzuTa Lin committed
996
            self.image = image
997
            self.filePath = unicodeFilePath
TzuTa Lin's avatar
TzuTa Lin committed
998
999
1000
1001
1002
1003
1004
            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()
1005
            self.addRecentFile(self.filePath)
TzuTa Lin's avatar
TzuTa Lin committed
1006
            self.toggleActions(True)
1007

1008
            # Label xml file and show bound box according to its filename
Wang Yinghao's avatar
Wang Yinghao committed
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
            # if self.usingPascalVocFormat is True:
            if self.defaultSaveDir is not None:
                basename = os.path.basename(
                    os.path.splitext(self.filePath)[0])
                xmlPath = os.path.join(self.defaultSaveDir, basename + XML_EXT)
                txtPath = os.path.join(self.defaultSaveDir, basename + TXT_EXT)

                """Annotation file priority:
                PascalXML > YOLO
                """
                if os.path.isfile(xmlPath):
1020
                    self.loadPascalXMLByFilename(xmlPath)
Wang Yinghao's avatar
Wang Yinghao committed
1021
1022
1023
1024
1025
1026
1027
1028
1029
                elif os.path.isfile(txtPath):
                    self.loadYOLOTXTByFilename(txtPath)
            else:
                xmlPath = os.path.splitext(filePath)[0] + XML_EXT
                txtPath = os.path.splitext(filePath)[0] + TXT_EXT
                if os.path.isfile(xmlPath):
                    self.loadPascalXMLByFilename(xmlPath)
                elif os.path.isfile(txtPath):
                    self.loadYOLOTXTByFilename(txtPath)
1030

1031
            self.setWindowTitle(__appname__ + ' ' + filePath)
1032
1033

            # Default : select last item if there is at least one item
1034
1035
            if self.labelList.count():
                self.labelList.setCurrentItem(self.labelList.item(self.labelList.count()-1))
1036
                self.labelList.item(self.labelList.count()-1).setSelected(True)
1037

Thibaut Mattio's avatar
Thibaut Mattio committed
1038
            self.canvas.setFocus(True)
TzuTa Lin's avatar
TzuTa Lin committed
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055