app.py 67.9 KB
Newer Older
1 2
# -*- coding: utf-8 -*-

K
Kentaro Wada 已提交
3
import functools
K
Kentaro Wada 已提交
4 5
import os
import os.path as osp
6
import re
K
Kentaro Wada 已提交
7
import webbrowser
8

9
import imgviz
10 11 12 13
from qtpy import QtCore
from qtpy.QtCore import Qt
from qtpy import QtGui
from qtpy import QtWidgets
M
Michael Pitidis 已提交
14

K
Kentaro Wada 已提交
15
from labelme import __appname__
K
Kentaro Wada 已提交
16
from labelme import PY2
K
Kentaro Wada 已提交
17
from labelme import QT5
K
Kentaro Wada 已提交
18

19
from . import utils
K
Kentaro Wada 已提交
20 21 22
from labelme.config import get_config
from labelme.label_file import LabelFile
from labelme.label_file import LabelFileError
23
from labelme.logger import logger
K
Kentaro Wada 已提交
24
from labelme.shape import Shape
25
from labelme.widgets import BrightnessContrastDialog
K
Kentaro Wada 已提交
26 27
from labelme.widgets import Canvas
from labelme.widgets import LabelDialog
28 29
from labelme.widgets import LabelListWidget
from labelme.widgets import LabelListWidgetItem
K
Kentaro Wada 已提交
30
from labelme.widgets import ToolBar
31
from labelme.widgets import UniqueLabelQListWidget
K
Kentaro Wada 已提交
32
from labelme.widgets import ZoomWidget
M
Michael Pitidis 已提交
33

34

35
# FIXME
36
# - [medium] Set max zoom value to something big enough for FitWidth/Window
37

K
Kentaro Wada 已提交
38
# TODO(unknown):
39
# - [high] Add polygon movement with arrow keys
40
# - [high] Deselect shape when clicking and already selected(?)
41
# - [low,maybe] Preview images on file dialogs.
42 43
# - Zoom is too "steppy".

M
Michael Pitidis 已提交
44

K
Kentaro Wada 已提交
45
LABEL_COLORMAP = imgviz.label_colormap(value=200)
K
Kentaro Wada 已提交
46 47 48 49
EXTENSIONS = [
    ".%s" % fmt.data().decode("ascii").lower()
    for fmt in QtGui.QImageReader.supportedImageFormats()
]
50 51


K
Kentaro Wada 已提交
52
class MainWindow(QtWidgets.QMainWindow):
M
Michael Pitidis 已提交
53

K
Kentaro Wada 已提交
54
    FIT_WINDOW, FIT_WIDTH, MANUAL_ZOOM = 0, 1, 2
M
Michael Pitidis 已提交
55

56 57 58 59 60 61 62
    def __init__(
        self,
        config=None,
        filename=None,
        output=None,
        output_file=None,
        output_dir=None,
K
Kentaro Wada 已提交
63
    ):
64
        if output is not None:
K
Kentaro Wada 已提交
65
            logger.warning(
K
Kentaro Wada 已提交
66
                "argument output is deprecated, use output_file instead"
67 68 69 70
            )
            if output_file is None:
                output_file = output

71
        # see labelme/config/default_config.yaml for valid configuration
K
Kentaro Wada 已提交
72 73
        if config is None:
            config = get_config()
74 75 76 77
        self._config = config

        super(MainWindow, self).__init__()
        self.setWindowTitle(__appname__)
K
Kentaro Wada 已提交
78

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

82 83
        self._noSelectionSlot = False

84
        # Main widgets and related state.
85 86
        self.labelDialog = LabelDialog(
            parent=self,
K
Kentaro Wada 已提交
87 88 89 90 91 92
            labels=self._config["labels"],
            sort_labels=self._config["sort_labels"],
            show_text_field=self._config["show_label_text_field"],
            completion=self._config["label_completion"],
            fit_to_content=self._config["fit_to_content"],
            flags=self._config["label_flags"],
93
        )
94

95
        self.labelList = LabelListWidget()
D
Douglas Long 已提交
96
        self.lastOpenDir = None
97

K
Kentaro Wada 已提交
98
        self.flag_dock = self.flag_widget = None
K
Kentaro Wada 已提交
99 100
        self.flag_dock = QtWidgets.QDockWidget(self.tr("Flags"), self)
        self.flag_dock.setObjectName("Flags")
K
Kentaro Wada 已提交
101
        self.flag_widget = QtWidgets.QListWidget()
K
Kentaro Wada 已提交
102 103
        if config["flags"]:
            self.loadFlags({k: False for k in config["flags"]})
K
Kentaro Wada 已提交
104 105 106
        self.flag_dock.setWidget(self.flag_widget)
        self.flag_widget.itemChanged.connect(self.setDirty)

107 108 109
        self.labelList.itemSelectionChanged.connect(self.labelSelectionChanged)
        self.labelList.itemDoubleClicked.connect(self.editLabel)
        self.labelList.itemChanged.connect(self.labelItemChanged)
110
        self.labelList.itemDropped.connect(self.labelOrderChanged)
C
Chen Zhang 已提交
111
        self.shape_dock = QtWidgets.QDockWidget(
K
Kentaro Wada 已提交
112
            self.tr("Polygon Labels"), self
C
Chen Zhang 已提交
113
        )
K
Kentaro Wada 已提交
114
        self.shape_dock.setObjectName("Labels")
115
        self.shape_dock.setWidget(self.labelList)
116

117
        self.uniqLabelList = UniqueLabelQListWidget()
K
Kentaro Wada 已提交
118 119 120 121 122 123 124 125
        self.uniqLabelList.setToolTip(
            self.tr(
                "Select label to start annotating for it. "
                "Press 'Esc' to deselect."
            )
        )
        if self._config["labels"]:
            for label in self._config["labels"]:
126 127 128 129
                item = self.uniqLabelList.createItemFromLabel(label)
                self.uniqLabelList.addItem(item)
                rgb = self._get_rgb_by_label(label)
                self.uniqLabelList.setItemLabel(item, label, rgb)
K
Kentaro Wada 已提交
130 131
        self.label_dock = QtWidgets.QDockWidget(self.tr(u"Label List"), self)
        self.label_dock.setObjectName(u"Label List")
132
        self.label_dock.setWidget(self.uniqLabelList)
M
Michael Pitidis 已提交
133

K
Kentaro Wada 已提交
134
        self.fileSearch = QtWidgets.QLineEdit()
K
Kentaro Wada 已提交
135
        self.fileSearch.setPlaceholderText(self.tr("Search Filename"))
K
Kentaro Wada 已提交
136
        self.fileSearch.textChanged.connect(self.fileSearchChanged)
K
Kentaro Wada 已提交
137 138
        self.fileListWidget = QtWidgets.QListWidget()
        self.fileListWidget.itemSelectionChanged.connect(
139 140
            self.fileSelectionChanged
        )
K
Kentaro Wada 已提交
141 142 143 144 145
        fileListLayout = QtWidgets.QVBoxLayout()
        fileListLayout.setContentsMargins(0, 0, 0, 0)
        fileListLayout.setSpacing(0)
        fileListLayout.addWidget(self.fileSearch)
        fileListLayout.addWidget(self.fileListWidget)
K
Kentaro Wada 已提交
146 147
        self.file_dock = QtWidgets.QDockWidget(self.tr(u"File List"), self)
        self.file_dock.setObjectName(u"Files")
K
Kentaro Wada 已提交
148 149
        fileListWidget = QtWidgets.QWidget()
        fileListWidget.setLayout(fileListLayout)
150
        self.file_dock.setWidget(fileListWidget)
D
Douglas Long 已提交
151

M
Michael Pitidis 已提交
152
        self.zoomWidget = ZoomWidget()
A
Aleksi J 已提交
153
        self.setAcceptDrops(True)
M
Michael Pitidis 已提交
154

155
        self.canvas = self.labelList.canvas = Canvas(
K
Kentaro Wada 已提交
156 157
            epsilon=self._config["epsilon"],
            double_click=self._config["canvas"]["double_click"],
158
        )
159
        self.canvas.zoomRequest.connect(self.zoomRequest)
160

161 162 163
        scrollArea = QtWidgets.QScrollArea()
        scrollArea.setWidget(self.canvas)
        scrollArea.setWidgetResizable(True)
164
        self.scrollBars = {
165 166
            Qt.Vertical: scrollArea.verticalScrollBar(),
            Qt.Horizontal: scrollArea.horizontalScrollBar(),
K
Kentaro Wada 已提交
167
        }
168
        self.canvas.scrollRequest.connect(self.scrollRequest)
M
Michael Pitidis 已提交
169

170
        self.canvas.newShape.connect(self.newShape)
171
        self.canvas.shapeMoved.connect(self.setDirty)
172
        self.canvas.selectionChanged.connect(self.shapeSelectionChanged)
M
Michael Pitidis 已提交
173
        self.canvas.drawingPolygon.connect(self.toggleDrawingSensitive)
174

175
        self.setCentralWidget(scrollArea)
D
Douglas Long 已提交
176

177
        features = QtWidgets.QDockWidget.DockWidgetFeatures()
K
Kentaro Wada 已提交
178 179
        for dock in ["flag_dock", "label_dock", "shape_dock", "file_dock"]:
            if self._config[dock]["closable"]:
180
                features = features | QtWidgets.QDockWidget.DockWidgetClosable
K
Kentaro Wada 已提交
181
            if self._config[dock]["floatable"]:
182
                features = features | QtWidgets.QDockWidget.DockWidgetFloatable
K
Kentaro Wada 已提交
183
            if self._config[dock]["movable"]:
184 185
                features = features | QtWidgets.QDockWidget.DockWidgetMovable
            getattr(self, dock).setFeatures(features)
K
Kentaro Wada 已提交
186
            if self._config[dock]["show"] is False:
187 188
                getattr(self, dock).setVisible(False)

K
Kentaro Wada 已提交
189
        self.addDockWidget(Qt.RightDockWidgetArea, self.flag_dock)
190 191 192
        self.addDockWidget(Qt.RightDockWidgetArea, self.label_dock)
        self.addDockWidget(Qt.RightDockWidgetArea, self.shape_dock)
        self.addDockWidget(Qt.RightDockWidgetArea, self.file_dock)
M
Michael Pitidis 已提交
193 194

        # Actions
195
        action = functools.partial(utils.newAction, self)
K
Kentaro Wada 已提交
196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217
        shortcuts = self._config["shortcuts"]
        quit = action(
            self.tr("&Quit"),
            self.close,
            shortcuts["quit"],
            "quit",
            self.tr("Quit application"),
        )
        open_ = action(
            self.tr("&Open"),
            self.openFile,
            shortcuts["open"],
            "open",
            self.tr("Open image or label file"),
        )
        opendir = action(
            self.tr("&Open Dir"),
            self.openDirDialog,
            shortcuts["open_dir"],
            "open",
            self.tr(u"Open Dir"),
        )
218
        openNextImg = action(
K
Kentaro Wada 已提交
219
            self.tr("&Next Image"),
220
            self.openNextImg,
K
Kentaro Wada 已提交
221 222 223
            shortcuts["open_next"],
            "next",
            self.tr(u"Open next (hold Ctl+Shift to copy labels)"),
224 225 226
            enabled=False,
        )
        openPrevImg = action(
K
Kentaro Wada 已提交
227
            self.tr("&Prev Image"),
228
            self.openPrevImg,
K
Kentaro Wada 已提交
229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247
            shortcuts["open_prev"],
            "prev",
            self.tr(u"Open prev (hold Ctl+Shift to copy labels)"),
            enabled=False,
        )
        save = action(
            self.tr("&Save"),
            self.saveFile,
            shortcuts["save"],
            "save",
            self.tr("Save labels to file"),
            enabled=False,
        )
        saveAs = action(
            self.tr("&Save As"),
            self.saveFileAs,
            shortcuts["save_as"],
            "save-as",
            self.tr("Save labels to a different file"),
248 249
            enabled=False,
        )
なるみ 已提交
250 251

        deleteFile = action(
K
Kentaro Wada 已提交
252
            self.tr("&Delete File"),
なるみ 已提交
253
            self.deleteFile,
K
Kentaro Wada 已提交
254 255 256 257 258
            shortcuts["delete_file"],
            "delete",
            self.tr("Delete current label file"),
            enabled=False,
        )
なるみ 已提交
259

260
        changeOutputDir = action(
K
Kentaro Wada 已提交
261
            self.tr("&Change Output Dir"),
262
            slot=self.changeOutputDirDialog,
K
Kentaro Wada 已提交
263 264 265
            shortcut=shortcuts["save_to"],
            icon="open",
            tip=self.tr(u"Change where annotations are loaded/saved"),
266
        )
K
Kentaro Wada 已提交
267 268

        saveAuto = action(
K
Kentaro Wada 已提交
269
            text=self.tr("Save &Automatically"),
K
Kentaro Wada 已提交
270
            slot=lambda x: self.actions.saveAuto.setChecked(x),
K
Kentaro Wada 已提交
271 272
            icon="save",
            tip=self.tr("Save automatically"),
K
Kentaro Wada 已提交
273 274 275
            checkable=True,
            enabled=True,
        )
K
Kentaro Wada 已提交
276
        saveAuto.setChecked(self._config["auto_save"])
K
Kentaro Wada 已提交
277

278
        saveWithImageData = action(
K
Kentaro Wada 已提交
279
            text="Save With Image Data",
K
Kentaro Wada 已提交
280
            slot=self.enableSaveImageWithData,
K
Kentaro Wada 已提交
281
            tip="Save image data in label file",
282
            checkable=True,
K
Kentaro Wada 已提交
283
            checked=self._config["store_data"],
284 285
        )

K
Kentaro Wada 已提交
286 287 288 289 290 291 292
        close = action(
            "&Close",
            self.closeFile,
            shortcuts["close"],
            "close",
            "Close current file",
        )
293

294
        toggle_keep_prev_mode = action(
K
Kentaro Wada 已提交
295
            self.tr("Keep Previous Annotation"),
296
            self.toggleKeepPrevMode,
K
Kentaro Wada 已提交
297 298
            shortcuts["toggle_keep_prev_mode"],
            None,
C
Chen Zhang 已提交
299
            self.tr('Toggle "keep pevious annotation" mode'),
K
Kentaro Wada 已提交
300 301 302
            checkable=True,
        )
        toggle_keep_prev_mode.setChecked(self._config["keep_prev"])
303

K
Kentaro Wada 已提交
304
        createMode = action(
K
Kentaro Wada 已提交
305 306 307 308 309
            self.tr("Create Polygons"),
            lambda: self.toggleDrawMode(False, createMode="polygon"),
            shortcuts["create_polygon"],
            "objects",
            self.tr("Start drawing polygons"),
310
            enabled=False,
K
Kentaro Wada 已提交
311
        )
K
Kentaro Wada 已提交
312
        createRectangleMode = action(
K
Kentaro Wada 已提交
313 314 315 316 317
            self.tr("Create Rectangle"),
            lambda: self.toggleDrawMode(False, createMode="rectangle"),
            shortcuts["create_rectangle"],
            "objects",
            self.tr("Start drawing rectangles"),
318
            enabled=False,
K
Kentaro Wada 已提交
319
        )
T
Tatiana Malygina 已提交
320
        createCircleMode = action(
K
Kentaro Wada 已提交
321 322 323 324 325
            self.tr("Create Circle"),
            lambda: self.toggleDrawMode(False, createMode="circle"),
            shortcuts["create_circle"],
            "objects",
            self.tr("Start drawing circles"),
326
            enabled=False,
T
Tatiana Malygina 已提交
327
        )
K
Kentaro Wada 已提交
328
        createLineMode = action(
K
Kentaro Wada 已提交
329 330 331 332 333
            self.tr("Create Line"),
            lambda: self.toggleDrawMode(False, createMode="line"),
            shortcuts["create_line"],
            "objects",
            self.tr("Start drawing lines"),
334
            enabled=False,
K
Kentaro Wada 已提交
335
        )
K
Kentaro Wada 已提交
336
        createPointMode = action(
K
Kentaro Wada 已提交
337 338 339 340 341
            self.tr("Create Point"),
            lambda: self.toggleDrawMode(False, createMode="point"),
            shortcuts["create_point"],
            "objects",
            self.tr("Start drawing points"),
342
            enabled=False,
K
Kentaro Wada 已提交
343
        )
S
Shohei Fujii 已提交
344
        createLineStripMode = action(
K
Kentaro Wada 已提交
345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382
            self.tr("Create LineStrip"),
            lambda: self.toggleDrawMode(False, createMode="linestrip"),
            shortcuts["create_linestrip"],
            "objects",
            self.tr("Start drawing linestrip. Ctrl+LeftClick ends creation."),
            enabled=False,
        )
        editMode = action(
            self.tr("Edit Polygons"),
            self.setEditMode,
            shortcuts["edit_polygon"],
            "edit",
            self.tr("Move and edit the selected polygons"),
            enabled=False,
        )

        delete = action(
            self.tr("Delete Polygons"),
            self.deleteSelectedShape,
            shortcuts["delete_polygon"],
            "cancel",
            self.tr("Delete the selected polygons"),
            enabled=False,
        )
        copy = action(
            self.tr("Duplicate Polygons"),
            self.copySelectedShape,
            shortcuts["duplicate_polygon"],
            "copy",
            self.tr("Create a duplicate of the selected polygons"),
            enabled=False,
        )
        undoLastPoint = action(
            self.tr("Undo last point"),
            self.canvas.undoLastPoint,
            shortcuts["undo_last_point"],
            "undo",
            self.tr("Undo last drawn point"),
S
Shohei Fujii 已提交
383 384
            enabled=False,
        )
K
Kentaro Wada 已提交
385
        addPointToEdge = action(
K
Kentaro Wada 已提交
386
            text=self.tr("Add Point to Edge"),
387
            slot=self.canvas.addPointToEdge,
K
Kentaro Wada 已提交
388 389 390
            shortcut=shortcuts["add_point_to_edge"],
            icon="edit",
            tip=self.tr("Add point to the nearest edge"),
K
Kentaro Wada 已提交
391 392
            enabled=False,
        )
K
Kentaro Wada 已提交
393
        removePoint = action(
K
Kentaro Wada 已提交
394
            text="Remove Selected Point",
K
Kentaro Wada 已提交
395
            slot=self.canvas.removeSelectedPoint,
K
Kentaro Wada 已提交
396 397
            icon="edit",
            tip="Remove selected point from polygon",
K
Kentaro Wada 已提交
398 399
            enabled=False,
        )
400

K
Kentaro Wada 已提交
401 402 403 404 405 406 407 408
        undo = action(
            self.tr("Undo"),
            self.undoShapeEdit,
            shortcuts["undo"],
            "undo",
            self.tr("Undo last add and edit of shape"),
            enabled=False,
        )
K
Kentaro Wada 已提交
409

K
Kentaro Wada 已提交
410 411 412 413 414 415 416 417 418 419 420 421 422 423
        hideAll = action(
            self.tr("&Hide\nPolygons"),
            functools.partial(self.togglePolygons, False),
            icon="eye",
            tip=self.tr("Hide all polygons"),
            enabled=False,
        )
        showAll = action(
            self.tr("&Show\nPolygons"),
            functools.partial(self.togglePolygons, True),
            icon="eye",
            tip=self.tr("Show all polygons"),
            enabled=False,
        )
M
Michael Pitidis 已提交
424

K
Kentaro Wada 已提交
425 426 427 428 429 430
        help = action(
            self.tr("&Tutorial"),
            self.tutorial,
            icon="help",
            tip=self.tr("Show tutorial page"),
        )
431

K
Kentaro Wada 已提交
432
        zoom = QtWidgets.QWidgetAction(self)
M
Michael Pitidis 已提交
433
        zoom.setDefaultWidget(self.zoomWidget)
434
        self.zoomWidget.setWhatsThis(
C
Chen Zhang 已提交
435
            self.tr(
K
Kentaro Wada 已提交
436 437
                "Zoom in or out of the image. Also accessible with "
                "{} and {} from the canvas."
C
Chen Zhang 已提交
438
            ).format(
439
                utils.fmtShortcut(
K
Kentaro Wada 已提交
440
                    "{},{}".format(shortcuts["zoom_in"], shortcuts["zoom_out"])
441
                ),
C
Chen Zhang 已提交
442
                utils.fmtShortcut(self.tr("Ctrl+Wheel")),
443 444
            )
        )
445 446
        self.zoomWidget.setEnabled(False)

K
Kentaro Wada 已提交
447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 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 495 496
        zoomIn = action(
            self.tr("Zoom &In"),
            functools.partial(self.addZoom, 1.1),
            shortcuts["zoom_in"],
            "zoom-in",
            self.tr("Increase zoom level"),
            enabled=False,
        )
        zoomOut = action(
            self.tr("&Zoom Out"),
            functools.partial(self.addZoom, 0.9),
            shortcuts["zoom_out"],
            "zoom-out",
            self.tr("Decrease zoom level"),
            enabled=False,
        )
        zoomOrg = action(
            self.tr("&Original size"),
            functools.partial(self.setZoom, 100),
            shortcuts["zoom_to_original"],
            "zoom",
            self.tr("Zoom to original size"),
            enabled=False,
        )
        fitWindow = action(
            self.tr("&Fit Window"),
            self.setFitWindow,
            shortcuts["fit_window"],
            "fit-window",
            self.tr("Zoom follows window size"),
            checkable=True,
            enabled=False,
        )
        fitWidth = action(
            self.tr("Fit &Width"),
            self.setFitWidth,
            shortcuts["fit_width"],
            "fit-width",
            self.tr("Zoom follows window width"),
            checkable=True,
            enabled=False,
        )
        brightnessContrast = action(
            "&Brightness Contrast",
            self.brightnessContrast,
            None,
            "color",
            "Adjust brightness and contrast",
            enabled=False,
        )
497
        # Group zoom controls into a list for easier toggling.
K
Kentaro Wada 已提交
498 499 500 501 502 503 504 505
        zoomActions = (
            self.zoomWidget,
            zoomIn,
            zoomOut,
            zoomOrg,
            fitWindow,
            fitWidth,
        )
K
Kentaro Wada 已提交
506 507
        self.zoomMode = self.FIT_WINDOW
        fitWindow.setChecked(Qt.Checked)
M
Michael Pitidis 已提交
508 509 510
        self.scalers = {
            self.FIT_WINDOW: self.scaleFitWindow,
            self.FIT_WIDTH: self.scaleFitWidth,
511 512
            # Set to one to scale to 100% when loading files.
            self.MANUAL_ZOOM: lambda: 1,
M
Michael Pitidis 已提交
513
        }
M
Michael Pitidis 已提交
514

K
Kentaro Wada 已提交
515 516 517 518 519 520 521 522
        edit = action(
            self.tr("&Edit Label"),
            self.editLabel,
            shortcuts["edit_label"],
            "edit",
            self.tr("Modify the label of the selected polygon"),
            enabled=False,
        )
M
Michael Pitidis 已提交
523

524
        fill_drawing = action(
K
Kentaro Wada 已提交
525
            self.tr("Fill Drawing Polygon"),
K
Kentaro Wada 已提交
526
            self.canvas.setFillDrawing,
527
            None,
K
Kentaro Wada 已提交
528 529
            "color",
            self.tr("Fill polygon while drawing"),
530 531 532
            checkable=True,
            enabled=True,
        )
K
Kentaro Wada 已提交
533
        fill_drawing.trigger()
M
Michael Pitidis 已提交
534

535
        # Lavel list context menu.
K
Kentaro Wada 已提交
536
        labelMenu = QtWidgets.QMenu()
537
        utils.addActions(labelMenu, (edit, delete))
M
Michael Pitidis 已提交
538
        self.labelList.setContextMenuPolicy(Qt.CustomContextMenu)
K
Kentaro Wada 已提交
539
        self.labelList.customContextMenuRequested.connect(
K
Kentaro Wada 已提交
540 541
            self.popLabelListMenu
        )
M
Michael Pitidis 已提交
542

M
Michael Pitidis 已提交
543
        # Store actions for further handling.
544
        self.actions = utils.struct(
K
Kentaro Wada 已提交
545
            saveAuto=saveAuto,
546
            saveWithImageData=saveWithImageData,
547
            changeOutputDir=changeOutputDir,
K
Kentaro Wada 已提交
548 549 550 551
            save=save,
            saveAs=saveAs,
            open=open_,
            close=close,
なるみ 已提交
552
            deleteFile=deleteFile,
553
            toggleKeepPrevMode=toggle_keep_prev_mode,
K
Kentaro Wada 已提交
554 555 556 557 558 559 560 561 562
            delete=delete,
            edit=edit,
            copy=copy,
            undoLastPoint=undoLastPoint,
            undo=undo,
            addPointToEdge=addPointToEdge,
            removePoint=removePoint,
            createMode=createMode,
            editMode=editMode,
K
Kentaro Wada 已提交
563
            createRectangleMode=createRectangleMode,
T
Tatiana Malygina 已提交
564
            createCircleMode=createCircleMode,
K
Kentaro Wada 已提交
565
            createLineMode=createLineMode,
K
Kentaro Wada 已提交
566
            createPointMode=createPointMode,
S
Shohei Fujii 已提交
567
            createLineStripMode=createLineStripMode,
K
Kentaro Wada 已提交
568 569 570 571 572 573
            zoom=zoom,
            zoomIn=zoomIn,
            zoomOut=zoomOut,
            zoomOrg=zoomOrg,
            fitWindow=fitWindow,
            fitWidth=fitWidth,
574
            brightnessContrast=brightnessContrast,
K
Kentaro Wada 已提交
575
            zoomActions=zoomActions,
K
Kentaro Wada 已提交
576 577
            openNextImg=openNextImg,
            openPrevImg=openPrevImg,
K
Kentaro Wada 已提交
578
            fileMenuActions=(open_, opendir, save, saveAs, close, quit),
579
            tool=(),
K
Kentaro Wada 已提交
580 581 582 583 584 585 586 587 588 589 590 591 592
            # XXX: need to add some actions here to activate the shortcut
            editMenu=(
                edit,
                copy,
                delete,
                None,
                undo,
                undoLastPoint,
                None,
                addPointToEdge,
                None,
                toggle_keep_prev_mode,
            ),
K
Kentaro Wada 已提交
593
            # menu shown at right click
594
            menu=(
K
Kentaro Wada 已提交
595 596
                createMode,
                createRectangleMode,
T
Tatiana Malygina 已提交
597
                createCircleMode,
K
Kentaro Wada 已提交
598 599
                createLineMode,
                createPointMode,
S
Shohei Fujii 已提交
600
                createLineStripMode,
K
Kentaro Wada 已提交
601 602 603 604 605 606
                editMode,
                edit,
                copy,
                delete,
                undo,
                undoLastPoint,
K
Kentaro Wada 已提交
607
                addPointToEdge,
608
                removePoint,
K
Kentaro Wada 已提交
609
            ),
K
Kentaro Wada 已提交
610 611 612 613
            onLoadActive=(
                close,
                createMode,
                createRectangleMode,
T
Tatiana Malygina 已提交
614
                createCircleMode,
K
Kentaro Wada 已提交
615
                createLineMode,
K
Kentaro Wada 已提交
616
                createPointMode,
S
Shohei Fujii 已提交
617
                createLineStripMode,
K
Kentaro Wada 已提交
618
                editMode,
619
                brightnessContrast,
K
Kentaro Wada 已提交
620
            ),
K
Kentaro Wada 已提交
621 622
            onShapesPresent=(saveAs, hideAll, showAll),
        )
M
Michael Pitidis 已提交
623

624
        self.canvas.edgeSelected.connect(self.canvasShapeEdgeSelected)
625
        self.canvas.vertexSelected.connect(self.actions.removePoint.setEnabled)
626

627
        self.menus = utils.struct(
K
Kentaro Wada 已提交
628 629 630 631 632
            file=self.menu(self.tr("&File")),
            edit=self.menu(self.tr("&Edit")),
            view=self.menu(self.tr("&View")),
            help=self.menu(self.tr("&Help")),
            recentFiles=QtWidgets.QMenu(self.tr("Open &Recent")),
K
Kentaro Wada 已提交
633 634 635
            labelList=labelMenu,
        )

636 637 638 639 640 641 642 643 644 645 646 647
        utils.addActions(
            self.menus.file,
            (
                open_,
                openNextImg,
                openPrevImg,
                opendir,
                self.menus.recentFiles,
                save,
                saveAs,
                saveAuto,
                changeOutputDir,
648
                saveWithImageData,
649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675
                close,
                deleteFile,
                None,
                quit,
            ),
        )
        utils.addActions(self.menus.help, (help,))
        utils.addActions(
            self.menus.view,
            (
                self.flag_dock.toggleViewAction(),
                self.label_dock.toggleViewAction(),
                self.shape_dock.toggleViewAction(),
                self.file_dock.toggleViewAction(),
                None,
                fill_drawing,
                None,
                hideAll,
                showAll,
                None,
                zoomIn,
                zoomOut,
                zoomOrg,
                None,
                fitWindow,
                fitWidth,
                None,
676
                brightnessContrast,
677 678
            ),
        )
679

M
Michael Pitidis 已提交
680
        self.menus.file.aboutToShow.connect(self.updateFileMenu)
681

682
        # Custom context menu for the canvas widget:
683 684 685 686
        utils.addActions(self.canvas.menus[0], self.actions.menu)
        utils.addActions(
            self.canvas.menus[1],
            (
K
Kentaro Wada 已提交
687 688
                action("&Copy here", self.copyShape),
                action("&Move here", self.moveShape),
689 690
            ),
        )
691

K
Kentaro Wada 已提交
692
        self.tools = self.toolbar("Tools")
693
        # Menu buttons on Left
694
        self.actions.tool = (
695 696 697 698 699
            open_,
            opendir,
            openNextImg,
            openPrevImg,
            save,
なるみ 已提交
700
            deleteFile,
701 702 703 704 705 706
            None,
            createMode,
            editMode,
            copy,
            delete,
            undo,
707
            brightnessContrast,
708 709 710 711
            None,
            zoom,
            fitWidth,
        )
M
Michael Pitidis 已提交
712

K
Kentaro Wada 已提交
713
        self.statusBar().showMessage(self.tr("%s started.") % __appname__)
M
Michael Pitidis 已提交
714 715
        self.statusBar().show()

K
Kentaro Wada 已提交
716
        if output_file is not None and self._config["auto_save"]:
717
            logger.warn(
K
Kentaro Wada 已提交
718 719 720
                "If `auto_save` argument is True, `output_file` argument "
                "is ignored and output filename is automatically "
                "set as IMAGE_BASENAME.json."
721 722 723 724
            )
        self.output_file = output_file
        self.output_dir = output_dir

725
        # Application state.
K
Kentaro Wada 已提交
726
        self.image = QtGui.QImage()
727
        self.imagePath = None
M
Michael Pitidis 已提交
728
        self.recentFiles = []
729
        self.maxRecent = 7
K
Kentaro Wada 已提交
730
        self.otherData = None
731 732
        self.zoom_level = 100
        self.fit_window = False
K
Kentaro Wada 已提交
733
        self.zoom_values = {}  # key=filename, value=(zoom_mode, zoom_value)
734
        self.brightnessContrast_values = {}
K
Kentaro Wada 已提交
735 736 737 738
        self.scroll_values = {
            Qt.Horizontal: {},
            Qt.Vertical: {},
        }  # key=filename, value=scroll_value
739

K
Kentaro Wada 已提交
740
        if filename is not None and osp.isdir(filename):
741
            self.importDirImages(filename, load=False)
742 743 744
        else:
            self.filename = filename

K
Kentaro Wada 已提交
745 746
        if config["file_search"]:
            self.fileSearch.setText(config["file_search"])
747 748
            self.fileSearchChanged()

749
        # XXX: Could be completely declarative.
750
        # Restore application settings.
K
Kentaro Wada 已提交
751
        self.settings = QtCore.QSettings("labelme", "labelme")
752
        # FIXME: QSettings.value can return None on PyQt4
K
Kentaro Wada 已提交
753 754 755
        self.recentFiles = self.settings.value("recentFiles", []) or []
        size = self.settings.value("window/size", QtCore.QSize(600, 500))
        position = self.settings.value("window/position", QtCore.QPoint(0, 0))
756 757 758
        self.resize(size)
        self.move(position)
        # or simply:
K
Kentaro Wada 已提交
759 760
        # self.restoreGeometry(settings['window/geometry']
        self.restoreState(
K
Kentaro Wada 已提交
761 762
            self.settings.value("window/state", QtCore.QByteArray())
        )
763

M
Michael Pitidis 已提交
764
        # Populate the File menu dynamically.
765
        self.updateFileMenu()
K
Kentaro Wada 已提交
766 767
        # Since loading the file may take some time,
        # make sure it runs in the background.
K
Kentaro Wada 已提交
768
        if self.filename is not None:
K
Kentaro Wada 已提交
769
            self.queueEvent(functools.partial(self.loadFile, self.filename))
770

771
        # Callbacks:
M
Michael Pitidis 已提交
772
        self.zoomWidget.valueChanged.connect(self.paintCanvas)
773

774
        self.populateModeActions()
M
Michael Pitidis 已提交
775

K
Kentaro Wada 已提交
776 777
        # self.firstStart = True
        # if self.firstStart:
778 779
        #    QWhatsThis.enterWhatsThisMode()

K
Kentaro Wada 已提交
780 781 782
    def menu(self, title, actions=None):
        menu = self.menuBar().addMenu(title)
        if actions:
783
            utils.addActions(menu, actions)
K
Kentaro Wada 已提交
784 785 786 787
        return menu

    def toolbar(self, title, actions=None):
        toolbar = ToolBar(title)
K
Kentaro Wada 已提交
788
        toolbar.setObjectName("%sToolBar" % title)
K
Kentaro Wada 已提交
789 790 791
        # toolbar.setOrientation(Qt.Vertical)
        toolbar.setToolButtonStyle(Qt.ToolButtonTextUnderIcon)
        if actions:
792
            utils.addActions(toolbar, actions)
K
Kentaro Wada 已提交
793 794 795
        self.addToolBar(Qt.LeftToolBarArea, toolbar)
        return toolbar

K
Kentaro Wada 已提交
796
    # Support Functions
797

M
Michael Pitidis 已提交
798
    def noShapes(self):
799
        return not len(self.labelList)
M
Michael Pitidis 已提交
800

801
    def populateModeActions(self):
802
        tool, menu = self.actions.tool, self.actions.menu
M
Michael Pitidis 已提交
803
        self.tools.clear()
804
        utils.addActions(self.tools, tool)
805
        self.canvas.menus[0].clear()
806
        utils.addActions(self.canvas.menus[0], menu)
807
        self.menus.edit.clear()
K
Kentaro Wada 已提交
808 809 810
        actions = (
            self.actions.createMode,
            self.actions.createRectangleMode,
T
Tatiana Malygina 已提交
811
            self.actions.createCircleMode,
K
Kentaro Wada 已提交
812
            self.actions.createLineMode,
K
Kentaro Wada 已提交
813
            self.actions.createPointMode,
S
Shohei Fujii 已提交
814
            self.actions.createLineStripMode,
K
Kentaro Wada 已提交
815 816
            self.actions.editMode,
        )
817
        utils.addActions(self.menus.edit, actions + self.actions.editMenu)
M
Michael Pitidis 已提交
818

819
    def setDirty(self):
K
Kentaro Wada 已提交
820 821
        if self._config["auto_save"] or self.actions.saveAuto.isChecked():
            label_file = osp.splitext(self.imagePath)[0] + ".json"
822
            if self.output_dir:
823 824
                label_file_without_path = osp.basename(label_file)
                label_file = osp.join(self.output_dir, label_file_without_path)
K
Kentaro Wada 已提交
825 826
            self.saveLabels(label_file)
            return
827 828
        self.dirty = True
        self.actions.save.setEnabled(True)
K
Kentaro Wada 已提交
829
        self.actions.undo.setEnabled(self.canvas.isShapeRestorable)
830 831
        title = __appname__
        if self.filename is not None:
K
Kentaro Wada 已提交
832
            title = "{} - {}*".format(title, self.filename)
833
        self.setWindowTitle(title)
834 835 836 837

    def setClean(self):
        self.dirty = False
        self.actions.save.setEnabled(False)
838
        self.actions.createMode.setEnabled(True)
K
Kentaro Wada 已提交
839
        self.actions.createRectangleMode.setEnabled(True)
T
Tatiana Malygina 已提交
840
        self.actions.createCircleMode.setEnabled(True)
K
Kentaro Wada 已提交
841
        self.actions.createLineMode.setEnabled(True)
K
Kentaro Wada 已提交
842
        self.actions.createPointMode.setEnabled(True)
S
Shohei Fujii 已提交
843
        self.actions.createLineStripMode.setEnabled(True)
844 845
        title = __appname__
        if self.filename is not None:
K
Kentaro Wada 已提交
846
            title = "{} - {}".format(title, self.filename)
847
        self.setWindowTitle(title)
848

なるみ 已提交
849 850 851 852 853
        if self.hasLabelFile():
            self.actions.deleteFile.setEnabled(True)
        else:
            self.actions.deleteFile.setEnabled(False)

854 855 856 857
    def toggleActions(self, value=True):
        """Enable/Disable widgets which depend on an opened image."""
        for z in self.actions.zoomActions:
            z.setEnabled(value)
M
Michael Pitidis 已提交
858 859
        for action in self.actions.onLoadActive:
            action.setEnabled(value)
860

861 862 863 864 865
    def canvasShapeEdgeSelected(self, selected, shape):
        self.actions.addPointToEdge.setEnabled(
            selected and shape and shape.canAddPoint()
        )

866
    def queueEvent(self, function):
K
Kentaro Wada 已提交
867
        QtCore.QTimer.singleShot(0, function)
868

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

M
Michael Pitidis 已提交
872 873 874
    def resetState(self):
        self.labelList.clear()
        self.filename = None
K
Kentaro Wada 已提交
875
        self.imagePath = None
M
Michael Pitidis 已提交
876 877
        self.imageData = None
        self.labelFile = None
K
Kentaro Wada 已提交
878
        self.otherData = None
M
Michael Pitidis 已提交
879 880
        self.canvas.resetState()

M
Michael Pitidis 已提交
881 882 883 884 885 886
    def currentItem(self):
        items = self.labelList.selectedItems()
        if items:
            return items[0]
        return None

887
    def addRecentFile(self, filename):
M
Michael Pitidis 已提交
888 889 890 891 892 893
        if filename in self.recentFiles:
            self.recentFiles.remove(filename)
        elif len(self.recentFiles) >= self.maxRecent:
            self.recentFiles.pop()
        self.recentFiles.insert(0, filename)

K
Kentaro Wada 已提交
894 895
    # Callbacks

K
Kentaro Wada 已提交
896 897 898 899 900 901
    def undoShapeEdit(self):
        self.canvas.restoreShape()
        self.labelList.clear()
        self.loadShapes(self.canvas.shapes)
        self.actions.undo.setEnabled(self.canvas.isShapeRestorable)

902
    def tutorial(self):
K
Kentaro Wada 已提交
903
        url = "https://github.com/wkentaro/labelme/tree/master/examples/tutorial"  # NOQA
K
Kentaro Wada 已提交
904
        webbrowser.open(url)
905

M
Michael Pitidis 已提交
906
    def toggleDrawingSensitive(self, drawing=True):
K
Kentaro Wada 已提交
907 908 909 910
        """Toggle drawing sensitive.

        In the middle of drawing, toggling between modes should be disabled.
        """
M
Michael Pitidis 已提交
911
        self.actions.editMode.setEnabled(not drawing)
912
        self.actions.undoLastPoint.setEnabled(drawing)
K
Kentaro Wada 已提交
913
        self.actions.undo.setEnabled(not drawing)
914
        self.actions.delete.setEnabled(not drawing)
M
Michael Pitidis 已提交
915

K
Kentaro Wada 已提交
916
    def toggleDrawMode(self, edit=True, createMode="polygon"):
M
Michael Pitidis 已提交
917
        self.canvas.setEditing(edit)
K
Kentaro Wada 已提交
918
        self.canvas.createMode = createMode
K
Kentaro Wada 已提交
919 920 921
        if edit:
            self.actions.createMode.setEnabled(True)
            self.actions.createRectangleMode.setEnabled(True)
T
Tatiana Malygina 已提交
922
            self.actions.createCircleMode.setEnabled(True)
K
Kentaro Wada 已提交
923 924
            self.actions.createLineMode.setEnabled(True)
            self.actions.createPointMode.setEnabled(True)
S
Shohei Fujii 已提交
925
            self.actions.createLineStripMode.setEnabled(True)
K
Kentaro Wada 已提交
926
        else:
K
Kentaro Wada 已提交
927
            if createMode == "polygon":
K
Kentaro Wada 已提交
928 929
                self.actions.createMode.setEnabled(False)
                self.actions.createRectangleMode.setEnabled(True)
T
Tatiana Malygina 已提交
930
                self.actions.createCircleMode.setEnabled(True)
K
Kentaro Wada 已提交
931 932
                self.actions.createLineMode.setEnabled(True)
                self.actions.createPointMode.setEnabled(True)
S
Shohei Fujii 已提交
933
                self.actions.createLineStripMode.setEnabled(True)
K
Kentaro Wada 已提交
934
            elif createMode == "rectangle":
K
Kentaro Wada 已提交
935 936
                self.actions.createMode.setEnabled(True)
                self.actions.createRectangleMode.setEnabled(False)
T
Tatiana Malygina 已提交
937
                self.actions.createCircleMode.setEnabled(True)
K
Kentaro Wada 已提交
938 939
                self.actions.createLineMode.setEnabled(True)
                self.actions.createPointMode.setEnabled(True)
S
Shohei Fujii 已提交
940
                self.actions.createLineStripMode.setEnabled(True)
K
Kentaro Wada 已提交
941
            elif createMode == "line":
K
Kentaro Wada 已提交
942 943
                self.actions.createMode.setEnabled(True)
                self.actions.createRectangleMode.setEnabled(True)
T
Tatiana Malygina 已提交
944
                self.actions.createCircleMode.setEnabled(True)
K
Kentaro Wada 已提交
945 946
                self.actions.createLineMode.setEnabled(False)
                self.actions.createPointMode.setEnabled(True)
S
Shohei Fujii 已提交
947
                self.actions.createLineStripMode.setEnabled(True)
K
Kentaro Wada 已提交
948
            elif createMode == "point":
K
Kentaro Wada 已提交
949 950
                self.actions.createMode.setEnabled(True)
                self.actions.createRectangleMode.setEnabled(True)
T
Tatiana Malygina 已提交
951
                self.actions.createCircleMode.setEnabled(True)
K
Kentaro Wada 已提交
952 953
                self.actions.createLineMode.setEnabled(True)
                self.actions.createPointMode.setEnabled(False)
S
Shohei Fujii 已提交
954
                self.actions.createLineStripMode.setEnabled(True)
T
Tatiana Malygina 已提交
955 956 957 958 959 960
            elif createMode == "circle":
                self.actions.createMode.setEnabled(True)
                self.actions.createRectangleMode.setEnabled(True)
                self.actions.createCircleMode.setEnabled(False)
                self.actions.createLineMode.setEnabled(True)
                self.actions.createPointMode.setEnabled(True)
S
Shohei Fujii 已提交
961 962 963 964 965 966 967 968
                self.actions.createLineStripMode.setEnabled(True)
            elif createMode == "linestrip":
                self.actions.createMode.setEnabled(True)
                self.actions.createRectangleMode.setEnabled(True)
                self.actions.createCircleMode.setEnabled(True)
                self.actions.createLineMode.setEnabled(True)
                self.actions.createPointMode.setEnabled(True)
                self.actions.createLineStripMode.setEnabled(False)
K
Kentaro Wada 已提交
969
            else:
K
Kentaro Wada 已提交
970
                raise ValueError("Unsupported createMode: %s" % createMode)
M
Michael Pitidis 已提交
971 972 973 974
        self.actions.editMode.setEnabled(not edit)

    def setEditMode(self):
        self.toggleDrawMode(True)
975

M
Michael Pitidis 已提交
976 977
    def updateFileMenu(self):
        current = self.filename
K
Kentaro Wada 已提交
978

M
Michael Pitidis 已提交
979
        def exists(filename):
K
Kentaro Wada 已提交
980
            return osp.exists(str(filename))
K
Kentaro Wada 已提交
981

982
        menu = self.menus.recentFiles
M
Michael Pitidis 已提交
983
        menu.clear()
984
        files = [f for f in self.recentFiles if f != current and exists(f)]
985
        for i, f in enumerate(files):
K
Kentaro Wada 已提交
986
            icon = utils.newIcon("labels")
K
Kentaro Wada 已提交
987
            action = QtWidgets.QAction(
K
Kentaro Wada 已提交
988 989
                icon, "&%d %s" % (i + 1, QtCore.QFileInfo(f).fileName()), self
            )
K
Kentaro Wada 已提交
990
            action.triggered.connect(functools.partial(self.loadRecent, f))
991
            menu.addAction(action)
M
Michael Pitidis 已提交
992

M
Michael Pitidis 已提交
993 994 995
    def popLabelListMenu(self, point):
        self.menus.labelList.exec_(self.labelList.mapToGlobal(point))

996 997
    def validateLabel(self, label):
        # no validation
K
Kentaro Wada 已提交
998
        if self._config["validate_label"] is None:
999 1000 1001
            return True

        for i in range(self.uniqLabelList.count()):
1002
            label_i = self.uniqLabelList.item(i).data(Qt.UserRole)
K
Kentaro Wada 已提交
1003
            if self._config["validate_label"] in ["exact"]:
K
Kentaro Wada 已提交
1004
                if label_i == label:
1005 1006 1007
                    return True
        return False

1008 1009
    def editLabel(self, item=None):
        if item and not isinstance(item, LabelListWidgetItem):
K
Kentaro Wada 已提交
1010
            raise TypeError("item must be LabelListWidgetItem type")
1011

M
Michael Pitidis 已提交
1012
        if not self.canvas.editing():
1013
            return
1014 1015 1016 1017
        if not item:
            item = self.currentItem()
        if item is None:
            return
1018
        shape = item.shape()
1019 1020
        if shape is None:
            return
K
Kentaro Wada 已提交
1021 1022 1023
        text, flags, group_id = self.labelDialog.popUp(
            text=shape.label, flags=shape.flags, group_id=shape.group_id,
        )
K
Kentaro Wada 已提交
1024 1025
        if text is None:
            return
1026
        if not self.validateLabel(text):
C
Chen Zhang 已提交
1027
            self.errorMessage(
K
Kentaro Wada 已提交
1028 1029 1030 1031
                self.tr("Invalid label"),
                self.tr("Invalid label '{}' with validation type '{}'").format(
                    text, self._config["validate_label"]
                ),
C
Chen Zhang 已提交
1032
            )
1033
            return
1034
        shape.label = text
1035
        shape.flags = flags
K
Kentaro Wada 已提交
1036 1037 1038 1039
        shape.group_id = group_id
        if shape.group_id is None:
            item.setText(shape.label)
        else:
K
Kentaro Wada 已提交
1040
            item.setText("{} ({})".format(shape.label, shape.group_id))
1041
        self.setDirty()
1042 1043
        if not self.uniqLabelList.findItemsByLabel(shape.label):
            item = QtWidgets.QListWidgetItem()
J
Jonas 已提交
1044
            item.setData(Qt.UserRole, shape.label)
1045
            self.uniqLabelList.addItem(item)
M
Michael Pitidis 已提交
1046

K
Kentaro Wada 已提交
1047 1048
    def fileSearchChanged(self):
        self.importDirImages(
K
Kentaro Wada 已提交
1049
            self.lastOpenDir, pattern=self.fileSearch.text(), load=False,
K
Kentaro Wada 已提交
1050 1051
        )

1052 1053 1054 1055 1056 1057
    def fileSelectionChanged(self):
        items = self.fileListWidget.selectedItems()
        if not items:
            return
        item = items[0]

1058 1059 1060
        if not self.mayContinue():
            return

D
Douglas Long 已提交
1061 1062 1063 1064 1065 1066
        currIndex = self.imageList.index(str(item.text()))
        if currIndex < len(self.imageList):
            filename = self.imageList[currIndex]
            if filename:
                self.loadFile(filename)

1067
    # React to canvas signals.
I
IlyaOvodov 已提交
1068 1069 1070 1071 1072 1073 1074 1075
    def shapeSelectionChanged(self, selected_shapes):
        self._noSelectionSlot = True
        for shape in self.canvas.selectedShapes:
            shape.selected = False
        self.labelList.clearSelection()
        self.canvas.selectedShapes = selected_shapes
        for shape in self.canvas.selectedShapes:
            shape.selected = True
1076 1077
            item = self.labelList.findItemByShape(shape)
            self.labelList.selectItem(item)
1078
            self.labelList.scrollToItem(item)
I
IlyaOvodov 已提交
1079
        self._noSelectionSlot = False
K
Kentaro Wada 已提交
1080 1081 1082 1083
        n_selected = len(selected_shapes)
        self.actions.delete.setEnabled(n_selected)
        self.actions.copy.setEnabled(n_selected)
        self.actions.edit.setEnabled(n_selected == 1)
1084

1085
    def addLabel(self, shape):
K
Kentaro Wada 已提交
1086 1087 1088
        if shape.group_id is None:
            text = shape.label
        else:
K
Kentaro Wada 已提交
1089
            text = "{} ({})".format(shape.label, shape.group_id)
1090 1091
        label_list_item = LabelListWidgetItem(text, shape)
        self.labelList.addItem(label_list_item)
1092 1093 1094 1095 1096
        if not self.uniqLabelList.findItemsByLabel(shape.label):
            item = self.uniqLabelList.createItemFromLabel(shape.label)
            self.uniqLabelList.addItem(item)
            rgb = self._get_rgb_by_label(shape.label)
            self.uniqLabelList.setItemLabel(item, shape.label, rgb)
K
Kentaro Wada 已提交
1097
        self.labelDialog.addLabelHistory(shape.label)
M
Michael Pitidis 已提交
1098 1099
        for action in self.actions.onShapesPresent:
            action.setEnabled(True)
1100

1101 1102
        rgb = self._get_rgb_by_label(shape.label)
        if rgb is None:
1103
            return
1104

K
Kentaro Wada 已提交
1105
        r, g, b = rgb
1106
        label_list_item.setText(
K
Kentaro Wada 已提交
1107 1108 1109
            '{} <font color="#{:02x}{:02x}{:02x}">●</font>'.format(
                text, r, g, b
            )
1110
        )
K
Kentaro Wada 已提交
1111 1112
        shape.line_color = QtGui.QColor(r, g, b)
        shape.vertex_fill_color = QtGui.QColor(r, g, b)
1113
        shape.hvertex_fill_color = QtGui.QColor(255, 255, 255)
K
Kentaro Wada 已提交
1114
        shape.fill_color = QtGui.QColor(r, g, b, 128)
1115
        shape.select_line_color = QtGui.QColor(255, 255, 255)
K
Kentaro Wada 已提交
1116
        shape.select_fill_color = QtGui.QColor(r, g, b, 155)
1117 1118

    def _get_rgb_by_label(self, label):
K
Kentaro Wada 已提交
1119
        if self._config["shape_color"] == "auto":
1120 1121
            item = self.uniqLabelList.findItemsByLabel(label)[0]
            label_id = self.uniqLabelList.indexFromItem(item).row() + 1
K
Kentaro Wada 已提交
1122
            label_id += self._config["shift_auto_shape_color"]
1123
            return LABEL_COLORMAP[label_id % len(LABEL_COLORMAP)]
K
Kentaro Wada 已提交
1124 1125 1126 1127 1128 1129 1130 1131
        elif (
            self._config["shape_color"] == "manual"
            and self._config["label_colors"]
            and label in self._config["label_colors"]
        ):
            return self._config["label_colors"][label]
        elif self._config["default_shape_color"]:
            return self._config["default_shape_color"]
1132

K
Kentaro Wada 已提交
1133 1134
    def remLabels(self, shapes):
        for shape in shapes:
1135
            item = self.labelList.findItemByShape(shape)
K
Kentaro Wada 已提交
1136
            self.labelList.removeItem(item)
1137

1138
    def loadShapes(self, shapes, replace=True):
I
IlyaOvodov 已提交
1139
        self._noSelectionSlot = True
K
Kentaro Wada 已提交
1140 1141
        for shape in shapes:
            self.addLabel(shape)
I
IlyaOvodov 已提交
1142 1143
        self.labelList.clearSelection()
        self._noSelectionSlot = False
1144
        self.canvas.loadShapes(shapes, replace=replace)
K
Kentaro Wada 已提交
1145

M
Michael Pitidis 已提交
1146 1147
    def loadLabels(self, shapes):
        s = []
K
Kentaro Wada 已提交
1148
        for shape in shapes:
K
Kentaro Wada 已提交
1149 1150 1151 1152 1153 1154
            label = shape["label"]
            points = shape["points"]
            shape_type = shape["shape_type"]
            flags = shape["flags"]
            group_id = shape["group_id"]
            other_data = shape["other_data"]
K
Kentaro Wada 已提交
1155

K
Kentaro Wada 已提交
1156
            shape = Shape(
K
Kentaro Wada 已提交
1157
                label=label, shape_type=shape_type, group_id=group_id,
K
Kentaro Wada 已提交
1158
            )
M
Michael Pitidis 已提交
1159
            for x, y in points:
K
Kentaro Wada 已提交
1160
                shape.addPoint(QtCore.QPointF(x, y))
1161
            shape.close()
1162 1163

            default_flags = {}
K
Kentaro Wada 已提交
1164 1165
            if self._config["label_flags"]:
                for pattern, keys in self._config["label_flags"].items():
K
Kentaro Wada 已提交
1166 1167 1168
                    if re.match(pattern, label):
                        for key in keys:
                            default_flags[key] = False
1169 1170
            shape.flags = default_flags
            shape.flags.update(flags)
1171
            shape.other_data = other_data
1172 1173

            s.append(shape)
K
Kentaro Wada 已提交
1174
        self.loadShapes(s)
M
Michael Pitidis 已提交
1175

K
Kentaro Wada 已提交
1176 1177 1178 1179 1180 1181 1182 1183
    def loadFlags(self, flags):
        self.flag_widget.clear()
        for key, flag in flags.items():
            item = QtWidgets.QListWidgetItem(key)
            item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
            item.setCheckState(Qt.Checked if flag else Qt.Unchecked)
            self.flag_widget.addItem(item)

M
Michael Pitidis 已提交
1184
    def saveLabels(self, filename):
M
Michael Pitidis 已提交
1185
        lf = LabelFile()
K
Kentaro Wada 已提交
1186

1187
        def format_shape(s):
1188
            data = s.other_data.copy()
K
Kentaro Wada 已提交
1189 1190 1191 1192 1193 1194 1195 1196 1197
            data.update(
                dict(
                    label=s.label.encode("utf-8") if PY2 else s.label,
                    points=[(p.x(), p.y()) for p in s.points],
                    group_id=s.group_id,
                    shape_type=s.shape_type,
                    flags=s.flags,
                )
            )
1198
            return data
1199

1200
        shapes = [format_shape(item.shape()) for item in self.labelList]
K
Kentaro Wada 已提交
1201 1202 1203 1204 1205 1206
        flags = {}
        for i in range(self.flag_widget.count()):
            item = self.flag_widget.item(i)
            key = item.text()
            flag = item.checkState() == Qt.Checked
            flags[key] = flag
1207
        try:
K
Kentaro Wada 已提交
1208 1209
            imagePath = osp.relpath(self.imagePath, osp.dirname(filename))
            imageData = self.imageData if self._config["store_data"] else None
1210
            if osp.dirname(filename) and not osp.exists(osp.dirname(filename)):
1211
                os.makedirs(osp.dirname(filename))
K
Kentaro Wada 已提交
1212 1213 1214 1215 1216
            lf.save(
                filename=filename,
                shapes=shapes,
                imagePath=imagePath,
                imageData=imageData,
1217 1218
                imageHeight=self.image.height(),
                imageWidth=self.image.width(),
K
Kentaro Wada 已提交
1219 1220 1221
                otherData=self.otherData,
                flags=flags,
            )
1222
            self.labelFile = lf
1223
            items = self.fileListWidget.findItems(
1224
                self.imagePath, Qt.MatchExactly
1225 1226
            )
            if len(items) > 0:
K
Kentaro Wada 已提交
1227
                if len(items) != 1:
K
Kentaro Wada 已提交
1228
                    raise RuntimeError("There are duplicate files.")
1229
                items[0].setCheckState(Qt.Checked)
K
Kentaro Wada 已提交
1230 1231
            # disable allows next and previous image to proceed
            # self.filename = filename
1232
            return True
K
Kentaro Wada 已提交
1233
        except LabelFileError as e:
C
Chen Zhang 已提交
1234
            self.errorMessage(
K
Kentaro Wada 已提交
1235
                self.tr("Error saving label data"), self.tr("<b>%s</b>") % e
C
Chen Zhang 已提交
1236
            )
1237
            return False
M
Michael Pitidis 已提交
1238

H
Hussein 已提交
1239
    def copySelectedShape(self):
I
IlyaOvodov 已提交
1240 1241 1242 1243 1244
        added_shapes = self.canvas.copySelectedShapes()
        self.labelList.clearSelection()
        for shape in added_shapes:
            self.addLabel(shape)
        self.setDirty()
1245

1246
    def labelSelectionChanged(self):
I
IlyaOvodov 已提交
1247 1248 1249 1250 1251
        if self._noSelectionSlot:
            return
        if self.canvas.editing():
            selected_shapes = []
            for item in self.labelList.selectedItems():
1252
                selected_shapes.append(item.shape())
I
IlyaOvodov 已提交
1253 1254
            if selected_shapes:
                self.canvas.selectShapes(selected_shapes)
K
Kentaro Wada 已提交
1255 1256
            else:
                self.canvas.deSelectShape()
1257

1258
    def labelItemChanged(self, item):
1259
        shape = item.shape()
1260
        self.canvas.setShapeVisible(shape, item.checkState() == Qt.Checked)
1261

1262 1263 1264 1265
    def labelOrderChanged(self):
        self.setDirty()
        self.canvas.loadShapes([item.shape() for item in self.labelList])

K
Kentaro Wada 已提交
1266 1267
    # Callback functions:

1268
    def newShape(self):
1269 1270 1271 1272
        """Pop-up and give focus to the label editor.

        position MUST be in global coordinates.
        """
1273
        items = self.uniqLabelList.selectedItems()
T
taashi-s 已提交
1274
        text = None
1275 1276
        if items:
            text = items[0].data(Qt.UserRole)
1277
        flags = {}
1278
        group_id = None
K
Kentaro Wada 已提交
1279
        if self._config["display_label_popup"] or not text:
1280
            previous_text = self.labelDialog.edit.text()
1281
            text, flags, group_id = self.labelDialog.popUp(text)
1282 1283
            if not text:
                self.labelDialog.edit.setText(previous_text)
1284

1285
        if text and not self.validateLabel(text):
C
Chen Zhang 已提交
1286
            self.errorMessage(
K
Kentaro Wada 已提交
1287 1288 1289 1290
                self.tr("Invalid label"),
                self.tr("Invalid label '{}' with validation type '{}'").format(
                    text, self._config["validate_label"]
                ),
C
Chen Zhang 已提交
1291
            )
K
Kentaro Wada 已提交
1292
            text = ""
1293
        if text:
I
IlyaOvodov 已提交
1294
            self.labelList.clearSelection()
K
Kentaro Wada 已提交
1295
            shape = self.canvas.setLastLabel(text, flags)
1296
            shape.group_id = group_id
K
Kentaro Wada 已提交
1297
            self.addLabel(shape)
1298
            self.actions.editMode.setEnabled(True)
K
Kentaro Wada 已提交
1299 1300
            self.actions.undoLastPoint.setEnabled(False)
            self.actions.undo.setEnabled(True)
1301
            self.setDirty()
1302 1303 1304
        else:
            self.canvas.undoLastLine()
            self.canvas.shapesBackups.pop()
1305

1306
    def scrollRequest(self, delta, orientation):
K
Kentaro Wada 已提交
1307
        units = -delta * 0.1  # natural scroll
1308
        bar = self.scrollBars[orientation]
K
Kentaro Wada 已提交
1309 1310 1311 1312 1313 1314
        value = bar.value() + bar.singleStep() * units
        self.setScroll(orientation, value)

    def setScroll(self, orientation, value):
        self.scrollBars[orientation].setValue(value)
        self.scroll_values[orientation][self.filename] = value
1315

M
Michael Pitidis 已提交
1316 1317 1318 1319 1320
    def setZoom(self, value):
        self.actions.fitWidth.setChecked(False)
        self.actions.fitWindow.setChecked(False)
        self.zoomMode = self.MANUAL_ZOOM
        self.zoomWidget.setValue(value)
K
Kentaro Wada 已提交
1321
        self.zoom_values[self.filename] = (self.zoomMode, value)
M
Michael Pitidis 已提交
1322

1323 1324
    def addZoom(self, increment=1.1):
        self.setZoom(self.zoomWidget.value() * increment)
M
Michael Pitidis 已提交
1325

1326
    def zoomRequest(self, delta, pos):
1327
        canvas_width_old = self.canvas.width()
1328 1329 1330
        units = 1.1
        if delta < 0:
            units = 0.9
1331
        self.addZoom(units)
1332 1333 1334 1335 1336 1337 1338 1339

        canvas_width_new = self.canvas.width()
        if canvas_width_old != canvas_width_new:
            canvas_scale_factor = canvas_width_new / canvas_width_old

            x_shift = round(pos.x() * canvas_scale_factor) - pos.x()
            y_shift = round(pos.y() * canvas_scale_factor) - pos.y()

K
Kentaro Wada 已提交
1340 1341 1342 1343 1344
            self.setScroll(
                Qt.Horizontal,
                self.scrollBars[Qt.Horizontal].value() + x_shift,
            )
            self.setScroll(
K
Kentaro Wada 已提交
1345
                Qt.Vertical, self.scrollBars[Qt.Vertical].value() + y_shift,
K
Kentaro Wada 已提交
1346
            )
1347

1348
    def setFitWindow(self, value=True):
M
Michael Pitidis 已提交
1349 1350 1351 1352 1353 1354 1355 1356 1357 1358
        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()
1359

1360 1361 1362 1363
    def onNewBrightnessContrast(self, qimage):
        self.canvas.loadPixmap(
            QtGui.QPixmap.fromImage(qimage), clear_shapes=False
        )
1364

1365
    def brightnessContrast(self, value):
1366
        dialog = BrightnessContrastDialog(
1367 1368 1369
            utils.img_data_to_pil(self.imageData),
            self.onNewBrightnessContrast,
            parent=self,
K
Kentaro Wada 已提交
1370
        )
1371 1372 1373 1374
        brightness, contrast = self.brightnessContrast_values.get(
            self.filename, (None, None)
        )
        if brightness is not None:
1375
            dialog.slider_brightness.setValue(brightness)
1376
        if contrast is not None:
1377
            dialog.slider_brightness.setValue(contrast)
1378 1379
        dialog.exec_()

1380 1381 1382 1383
        brightness = dialog.slider_brightness.value()
        contrast = dialog.slider_brightness.value()
        self.brightnessContrast_values[self.filename] = (brightness, contrast)

M
Michael Pitidis 已提交
1384
    def togglePolygons(self, value):
1385
        for item in self.labelList:
M
Michael Pitidis 已提交
1386
            item.setCheckState(Qt.Checked if value else Qt.Unchecked)
1387

1388 1389
    def loadFile(self, filename=None):
        """Load the specified file, or the last opened file if None."""
K
Kentaro Wada 已提交
1390
        # changing fileListWidget loads file
K
Kentaro Wada 已提交
1391 1392 1393
        if filename in self.imageList and (
            self.fileListWidget.currentRow() != self.imageList.index(filename)
        ):
K
Kentaro Wada 已提交
1394
            self.fileListWidget.setCurrentRow(self.imageList.index(filename))
1395
            self.fileListWidget.repaint()
K
Kentaro Wada 已提交
1396 1397
            return

M
Michael Pitidis 已提交
1398 1399
        self.resetState()
        self.canvas.setEnabled(False)
1400
        if filename is None:
K
Kentaro Wada 已提交
1401
            filename = self.settings.value("filename", "")
K
Kentaro Wada 已提交
1402
        filename = str(filename)
K
Kentaro Wada 已提交
1403
        if not QtCore.QFile.exists(filename):
1404
            self.errorMessage(
K
Kentaro Wada 已提交
1405 1406
                self.tr("Error opening file"),
                self.tr("No such file: <b>%s</b>") % filename,
C
Chen Zhang 已提交
1407
            )
1408 1409
            return False
        # assumes same name, but json extension
C
Chen Zhang 已提交
1410
        self.status(self.tr("Loading %s...") % osp.basename(str(filename)))
K
Kentaro Wada 已提交
1411
        label_file = osp.splitext(filename)[0] + ".json"
1412
        if self.output_dir:
1413 1414
            label_file_without_path = osp.basename(label_file)
            label_file = osp.join(self.output_dir, label_file_without_path)
K
Kentaro Wada 已提交
1415 1416 1417
        if QtCore.QFile.exists(label_file) and LabelFile.is_label_file(
            label_file
        ):
1418 1419 1420
            try:
                self.labelFile = LabelFile(label_file)
            except LabelFileError as e:
K
Kentaro Wada 已提交
1421
                self.errorMessage(
K
Kentaro Wada 已提交
1422
                    self.tr("Error opening file"),
C
Chen Zhang 已提交
1423 1424 1425
                    self.tr(
                        "<p><b>%s</b></p>"
                        "<p>Make sure <i>%s</i> is a valid label file."
K
Kentaro Wada 已提交
1426 1427
                    )
                    % (e, label_file),
C
Chen Zhang 已提交
1428
                )
C
Chen Zhang 已提交
1429
                self.status(self.tr("Error reading %s") % label_file)
1430
                return False
1431
            self.imageData = self.labelFile.imageData
K
Kentaro Wada 已提交
1432
            self.imagePath = osp.join(
K
Kentaro Wada 已提交
1433
                osp.dirname(label_file), self.labelFile.imagePath,
K
Kentaro Wada 已提交
1434
            )
K
Kentaro Wada 已提交
1435
            self.otherData = self.labelFile.otherData
1436
        else:
1437
            self.imageData = LabelFile.load_image_file(filename)
1438
            if self.imageData:
1439 1440
                self.imagePath = filename
            self.labelFile = None
K
Kentaro Wada 已提交
1441
        image = QtGui.QImage.fromData(self.imageData)
1442

1443
        if image.isNull():
K
Kentaro Wada 已提交
1444 1445 1446 1447
            formats = [
                "*.{}".format(fmt.data().decode())
                for fmt in QtGui.QImageReader.supportedImageFormats()
            ]
1448
            self.errorMessage(
K
Kentaro Wada 已提交
1449
                self.tr("Error opening file"),
C
Chen Zhang 已提交
1450
                self.tr(
K
Kentaro Wada 已提交
1451 1452 1453
                    "<p>Make sure <i>{0}</i> is a valid image file.<br/>"
                    "Supported image formats: {1}</p>"
                ).format(filename, ",".join(formats)),
C
Chen Zhang 已提交
1454
            )
C
Chen Zhang 已提交
1455
            self.status(self.tr("Error reading %s") % filename)
1456 1457 1458
            return False
        self.image = image
        self.filename = filename
K
Kentaro Wada 已提交
1459
        if self._config["keep_prev"]:
1460
            prev_shapes = self.canvas.shapes
K
Kentaro Wada 已提交
1461
        self.canvas.loadPixmap(QtGui.QPixmap.fromImage(image))
K
Kentaro Wada 已提交
1462
        flags = {k: False for k in self._config["flags"] or []}
1463 1464
        if self.labelFile:
            self.loadLabels(self.labelFile.shapes)
K
Kentaro Wada 已提交
1465
            if self.labelFile.flags is not None:
1466 1467
                flags.update(self.labelFile.flags)
        self.loadFlags(flags)
K
Kentaro Wada 已提交
1468
        if self._config["keep_prev"] and self.noShapes():
1469
            self.loadShapes(prev_shapes, replace=False)
K
Kentaro Wada 已提交
1470 1471 1472
            self.setDirty()
        else:
            self.setClean()
1473
        self.canvas.setEnabled(True)
K
Kentaro Wada 已提交
1474
        # set zoom values
K
Kentaro Wada 已提交
1475
        is_initial_load = not self.zoom_values
1476
        if self.filename in self.zoom_values:
K
Kentaro Wada 已提交
1477 1478
            self.zoomMode = self.zoom_values[self.filename][0]
            self.setZoom(self.zoom_values[self.filename][1])
K
Kentaro Wada 已提交
1479
        elif is_initial_load or not self._config["keep_prev_scale"]:
1480
            self.adjustScale(initial=True)
K
Kentaro Wada 已提交
1481 1482 1483 1484 1485 1486
        # set scroll values
        for orientation in self.scroll_values:
            if self.filename in self.scroll_values[orientation]:
                self.setScroll(
                    orientation, self.scroll_values[orientation][self.filename]
                )
1487
        # set brightness constrast values
1488
        dialog = BrightnessContrastDialog(
1489 1490 1491
            utils.img_data_to_pil(self.imageData),
            self.onNewBrightnessContrast,
            parent=self,
1492 1493 1494 1495 1496 1497 1498
        )
        brightness, contrast = self.brightnessContrast_values.get(
            self.filename, (None, None)
        )
        if self._config["keep_prev_brightness"] and self.recentFiles:
            brightness, _ = self.brightnessContrast_values.get(
                self.recentFiles[0], (None, None)
1499
            )
1500 1501 1502 1503 1504
        if self._config["keep_prev_contrast"] and self.recentFiles:
            _, contrast = self.brightnessContrast_values.get(
                self.recentFiles[0], (None, None)
            )
        if brightness is not None:
1505
            dialog.slider_brightness.setValue(brightness)
1506
        if contrast is not None:
1507
            dialog.slider_brightness.setValue(contrast)
1508 1509
        self.brightnessContrast_values[self.filename] = (brightness, contrast)
        if brightness is not None or contrast is not None:
1510
            dialog.onNewValue(None)
1511 1512 1513
        self.paintCanvas()
        self.addRecentFile(self.filename)
        self.toggleActions(True)
C
Chen Zhang 已提交
1514
        self.status(self.tr("Loaded %s") % osp.basename(str(filename)))
1515
        return True
1516

M
Michael Pitidis 已提交
1517
    def resizeEvent(self, event):
K
Kentaro Wada 已提交
1518 1519 1520 1521 1522
        if (
            self.canvas
            and not self.image.isNull()
            and self.zoomMode != self.MANUAL_ZOOM
        ):
M
Michael Pitidis 已提交
1523
            self.adjustScale()
M
Michael Pitidis 已提交
1524 1525
        super(MainWindow, self).resizeEvent(event)

M
Michael Pitidis 已提交
1526 1527
    def paintCanvas(self):
        assert not self.image.isNull(), "cannot paint null image"
M
Michael Pitidis 已提交
1528
        self.canvas.scale = 0.01 * self.zoomWidget.value()
M
Michael Pitidis 已提交
1529
        self.canvas.adjustSize()
M
Michael Pitidis 已提交
1530
        self.canvas.update()
M
Michael Pitidis 已提交
1531

1532 1533
    def adjustScale(self, initial=False):
        value = self.scalers[self.FIT_WINDOW if initial else self.zoomMode]()
K
Kentaro Wada 已提交
1534 1535 1536
        value = int(100 * value)
        self.zoomWidget.setValue(value)
        self.zoom_values[self.filename] = (self.zoomMode, value)
M
Michael Pitidis 已提交
1537 1538

    def scaleFitWindow(self):
K
Kentaro Wada 已提交
1539 1540
        """Figure out the size of the pixmap to fit the main widget."""
        e = 2.0  # So that no scrollbars are generated.
M
Michael Pitidis 已提交
1541 1542
        w1 = self.centralWidget().width() - e
        h1 = self.centralWidget().height() - e
K
Kentaro Wada 已提交
1543
        a1 = w1 / h1
M
Michael Pitidis 已提交
1544 1545 1546 1547 1548 1549
        # 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

M
Michael Pitidis 已提交
1550 1551 1552 1553
    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()
M
Michael Pitidis 已提交
1554

K
Kentaro Wada 已提交
1555
    def enableSaveImageWithData(self, enabled):
K
Kentaro Wada 已提交
1556
        self._config["store_data"] = enabled
K
Kentaro Wada 已提交
1557 1558
        self.actions.saveWithImageData.setChecked(enabled)

1559
    def closeEvent(self, event):
1560 1561
        if not self.mayContinue():
            event.ignore()
K
Kentaro Wada 已提交
1562
        self.settings.setValue(
K
Kentaro Wada 已提交
1563 1564 1565 1566 1567 1568
            "filename", self.filename if self.filename else ""
        )
        self.settings.setValue("window/size", self.size())
        self.settings.setValue("window/position", self.pos())
        self.settings.setValue("window/state", self.saveState())
        self.settings.setValue("recentFiles", self.recentFiles)
H
Hussein 已提交
1569
        # ask the use for where to save the labels
K
Kentaro Wada 已提交
1570
        # self.settings.setValue('window/geometry', self.saveGeometry())
M
Michael Pitidis 已提交
1571

A
Aleksi J 已提交
1572 1573 1574 1575 1576 1577 1578 1579 1580 1581 1582 1583 1584 1585 1586
    def dragEnterEvent(self, event):
        if event.mimeData().hasUrls():
            items = [i.toLocalFile() for i in event.mimeData().urls()]
            if any([i.lower().endswith(tuple(EXTENSIONS)) for i in items]):
                event.accept()
        else:
            event.ignore()

    def dropEvent(self, event):
        if not self.mayContinue():
            event.ignore()
            return
        items = [i.toLocalFile() for i in event.mimeData().urls()]
        self.importDroppedImageFiles(items)

K
Kentaro Wada 已提交
1587
    # User Dialogs #
1588

1589 1590 1591 1592
    def loadRecent(self, filename):
        if self.mayContinue():
            self.loadFile(filename)

D
Douglas Long 已提交
1593
    def openPrevImg(self, _value=False):
K
Kentaro Wada 已提交
1594
        keep_prev = self._config["keep_prev"]
1595
        if Qt.KeyboardModifiers() == (Qt.ControlModifier | Qt.ShiftModifier):
K
Kentaro Wada 已提交
1596
            self._config["keep_prev"] = True
1597

D
Douglas Long 已提交
1598 1599 1600 1601 1602 1603 1604 1605 1606 1607 1608 1609 1610 1611 1612
        if not self.mayContinue():
            return

        if len(self.imageList) <= 0:
            return

        if self.filename is None:
            return

        currIndex = self.imageList.index(self.filename)
        if currIndex - 1 >= 0:
            filename = self.imageList[currIndex - 1]
            if filename:
                self.loadFile(filename)

K
Kentaro Wada 已提交
1613
        self._config["keep_prev"] = keep_prev
1614

K
Kentaro Wada 已提交
1615
    def openNextImg(self, _value=False, load=True):
K
Kentaro Wada 已提交
1616
        keep_prev = self._config["keep_prev"]
1617
        if Qt.KeyboardModifiers() == (Qt.ControlModifier | Qt.ShiftModifier):
K
Kentaro Wada 已提交
1618
            self._config["keep_prev"] = True
1619

D
Douglas Long 已提交
1620 1621 1622 1623 1624 1625 1626 1627 1628 1629 1630 1631 1632
        if not self.mayContinue():
            return

        if len(self.imageList) <= 0:
            return

        filename = None
        if self.filename is None:
            filename = self.imageList[0]
        else:
            currIndex = self.imageList.index(self.filename)
            if currIndex + 1 < len(self.imageList):
                filename = self.imageList[currIndex + 1]
K
Kentaro Wada 已提交
1633 1634
            else:
                filename = self.imageList[-1]
K
Kentaro Wada 已提交
1635
        self.filename = filename
D
Douglas Long 已提交
1636

K
Kentaro Wada 已提交
1637 1638
        if self.filename and load:
            self.loadFile(self.filename)
D
Douglas Long 已提交
1639

K
Kentaro Wada 已提交
1640
        self._config["keep_prev"] = keep_prev
1641

1642 1643
    def openFile(self, _value=False):
        if not self.mayContinue():
1644
            return
K
Kentaro Wada 已提交
1645 1646 1647 1648 1649 1650 1651 1652
        path = osp.dirname(str(self.filename)) if self.filename else "."
        formats = [
            "*.{}".format(fmt.data().decode())
            for fmt in QtGui.QImageReader.supportedImageFormats()
        ]
        filters = self.tr("Image & Label files (%s)") % " ".join(
            formats + ["*%s" % LabelFile.suffix]
        )
K
Kentaro Wada 已提交
1653
        filename = QtWidgets.QFileDialog.getOpenFileName(
K
Kentaro Wada 已提交
1654 1655 1656 1657 1658
            self,
            self.tr("%s - Choose Image or Label file") % __appname__,
            path,
            filters,
        )
1659
        if QT5:
K
Kentaro Wada 已提交
1660 1661
            filename, _ = filename
        filename = str(filename)
M
Michael Pitidis 已提交
1662
        if filename:
1663
            self.loadFile(filename)
M
Michael Pitidis 已提交
1664

1665 1666 1667 1668 1669 1670 1671 1672
    def changeOutputDirDialog(self, _value=False):
        default_output_dir = self.output_dir
        if default_output_dir is None and self.filename:
            default_output_dir = osp.dirname(self.filename)
        if default_output_dir is None:
            default_output_dir = self.currentPath()

        output_dir = QtWidgets.QFileDialog.getExistingDirectory(
C
Chen Zhang 已提交
1673
            self,
K
Kentaro Wada 已提交
1674
            self.tr("%s - Save/Load Annotations in Directory") % __appname__,
1675
            default_output_dir,
K
Kentaro Wada 已提交
1676 1677
            QtWidgets.QFileDialog.ShowDirsOnly
            | QtWidgets.QFileDialog.DontResolveSymlinks,
1678 1679 1680 1681 1682 1683 1684 1685 1686
        )
        output_dir = str(output_dir)

        if not output_dir:
            return

        self.output_dir = output_dir

        self.statusBar().showMessage(
K
Kentaro Wada 已提交
1687 1688 1689
            self.tr("%s . Annotations will be saved/loaded in %s")
            % ("Change Annotations Dir", self.output_dir)
        )
1690 1691 1692 1693 1694 1695 1696 1697
        self.statusBar().show()

        current_filename = self.filename
        self.importDirImages(self.lastOpenDir, load=False)

        if current_filename in self.imageList:
            # retain currently selected file
            self.fileListWidget.setCurrentRow(
K
Kentaro Wada 已提交
1698 1699
                self.imageList.index(current_filename)
            )
1700 1701
            self.fileListWidget.repaint()

1702
    def saveFile(self, _value=False):
M
Michael Pitidis 已提交
1703
        assert not self.image.isNull(), "cannot save empty image"
1704 1705 1706 1707 1708 1709 1710 1711
        if self.labelFile:
            # DL20180323 - overwrite when in directory
            self._saveFile(self.labelFile.filename)
        elif self.output_file:
            self._saveFile(self.output_file)
            self.close()
        else:
            self._saveFile(self.saveFileDialog())
M
Michael Pitidis 已提交
1712 1713 1714

    def saveFileAs(self, _value=False):
        assert not self.image.isNull(), "cannot save empty image"
1715
        self._saveFile(self.saveFileDialog())
M
Michael Pitidis 已提交
1716 1717

    def saveFileDialog(self):
K
Kentaro Wada 已提交
1718 1719
        caption = self.tr("%s - Choose File") % __appname__
        filters = self.tr("Label files (*%s)") % LabelFile.suffix
1720 1721 1722 1723 1724 1725 1726 1727
        if self.output_dir:
            dlg = QtWidgets.QFileDialog(
                self, caption, self.output_dir, filters
            )
        else:
            dlg = QtWidgets.QFileDialog(
                self, caption, self.currentPath(), filters
            )
1728
        dlg.setDefaultSuffix(LabelFile.suffix[1:])
K
Kentaro Wada 已提交
1729 1730 1731
        dlg.setAcceptMode(QtWidgets.QFileDialog.AcceptSave)
        dlg.setOption(QtWidgets.QFileDialog.DontConfirmOverwrite, False)
        dlg.setOption(QtWidgets.QFileDialog.DontUseNativeDialog, False)
1732
        basename = osp.basename(osp.splitext(self.filename)[0])
1733 1734 1735 1736 1737 1738 1739 1740
        if self.output_dir:
            default_labelfile_name = osp.join(
                self.output_dir, basename + LabelFile.suffix
            )
        else:
            default_labelfile_name = osp.join(
                self.currentPath(), basename + LabelFile.suffix
            )
1741
        filename = dlg.getSaveFileName(
K
Kentaro Wada 已提交
1742 1743 1744 1745 1746
            self,
            self.tr("Choose File"),
            default_labelfile_name,
            self.tr("Label files (*%s)") % LabelFile.suffix,
        )
1747
        if isinstance(filename, tuple):
K
Kentaro Wada 已提交
1748 1749
            filename, _ = filename
        return filename
M
Michael Pitidis 已提交
1750 1751 1752 1753 1754

    def _saveFile(self, filename):
        if filename and self.saveLabels(filename):
            self.addRecentFile(filename)
            self.setClean()
1755

M
Michael Pitidis 已提交
1756 1757 1758 1759
    def closeFile(self, _value=False):
        if not self.mayContinue():
            return
        self.resetState()
1760
        self.setClean()
1761
        self.toggleActions(False)
M
Michael Pitidis 已提交
1762
        self.canvas.setEnabled(False)
M
Michael Pitidis 已提交
1763
        self.actions.saveAs.setEnabled(False)
M
Michael Pitidis 已提交
1764

なるみ 已提交
1765
    def getLabelFile(self):
K
Kentaro Wada 已提交
1766
        if self.filename.lower().endswith(".json"):
なるみ 已提交
1767
            label_file = self.filename
なるみ 已提交
1768
        else:
K
Kentaro Wada 已提交
1769
            label_file = osp.splitext(self.filename)[0] + ".json"
なるみ 已提交
1770 1771 1772 1773

        return label_file

    def deleteFile(self):
なるみ 已提交
1774
        mb = QtWidgets.QMessageBox
K
Kentaro Wada 已提交
1775 1776 1777 1778 1779
        msg = self.tr(
            "You are about to permanently delete this label file, "
            "proceed anyway?"
        )
        answer = mb.warning(self, self.tr("Attention"), msg, mb.Yes | mb.No)
K
Kentaro Wada 已提交
1780 1781 1782 1783 1784 1785
        if answer != mb.Yes:
            return

        label_file = self.getLabelFile()
        if osp.exists(label_file):
            os.remove(label_file)
K
Kentaro Wada 已提交
1786
            logger.info("Label file is removed: {}".format(label_file))
なるみ 已提交
1787

K
Kentaro Wada 已提交
1788 1789
            item = self.fileListWidget.currentItem()
            item.setCheckState(Qt.Unchecked)
なるみ 已提交
1790

1791
            self.resetState()
なるみ 已提交
1792

1793
    # Message Dialogs. #
M
Michael Pitidis 已提交
1794
    def hasLabels(self):
1795
        if self.noShapes():
K
Kentaro Wada 已提交
1796
            self.errorMessage(
K
Kentaro Wada 已提交
1797 1798 1799
                "No objects labeled",
                "You must label at least one object to save the file.",
            )
M
Michael Pitidis 已提交
1800 1801 1802
            return False
        return True

なるみ 已提交
1803 1804 1805 1806 1807 1808 1809
    def hasLabelFile(self):
        if self.filename is None:
            return False

        label_file = self.getLabelFile()
        return osp.exists(label_file)

1810
    def mayContinue(self):
1811 1812
        if not self.dirty:
            return True
K
Kentaro Wada 已提交
1813
        mb = QtWidgets.QMessageBox
C
Chen Zhang 已提交
1814
        msg = self.tr('Save annotations to "{}" before closing?').format(
K
Kentaro Wada 已提交
1815 1816 1817 1818 1819 1820 1821 1822 1823
            self.filename
        )
        answer = mb.question(
            self,
            self.tr("Save annotations?"),
            msg,
            mb.Save | mb.Discard | mb.Cancel,
            mb.Save,
        )
1824 1825 1826 1827 1828
        if answer == mb.Discard:
            return True
        elif answer == mb.Save:
            self.saveFile()
            return True
K
Kentaro Wada 已提交
1829
        else:  # answer == mb.Cancel
1830
            return False
M
Michael Pitidis 已提交
1831

1832
    def errorMessage(self, title, message):
K
Kentaro Wada 已提交
1833
        return QtWidgets.QMessageBox.critical(
K
Kentaro Wada 已提交
1834 1835
            self, title, "<p><b>%s</b></p>%s" % (title, message)
        )
1836

1837
    def currentPath(self):
K
Kentaro Wada 已提交
1838
        return osp.dirname(str(self.filename)) if self.filename else "."
M
Michael Pitidis 已提交
1839

1840
    def toggleKeepPrevMode(self):
K
Kentaro Wada 已提交
1841
        self._config["keep_prev"] = not self._config["keep_prev"]
1842

1843
    def deleteSelectedShape(self):
K
Kentaro Wada 已提交
1844
        yes, no = QtWidgets.QMessageBox.Yes, QtWidgets.QMessageBox.No
C
Chen Zhang 已提交
1845
        msg = self.tr(
K
Kentaro Wada 已提交
1846 1847
            "You are about to permanently delete {} polygons, "
            "proceed anyway?"
C
Chen Zhang 已提交
1848 1849
        ).format(len(self.canvas.selectedShapes))
        if yes == QtWidgets.QMessageBox.warning(
K
Kentaro Wada 已提交
1850 1851
            self, self.tr("Attention"), msg, yes | no, yes
        ):
I
IlyaOvodov 已提交
1852
            self.remLabels(self.canvas.deleteSelected())
1853
            self.setDirty()
M
Michael Pitidis 已提交
1854 1855 1856
            if self.noShapes():
                for action in self.actions.onShapesPresent:
                    action.setEnabled(False)
M
Martijn Buijs 已提交
1857

1858 1859
    def copyShape(self):
        self.canvas.endMove(copy=True)
I
IlyaOvodov 已提交
1860 1861 1862
        self.labelList.clearSelection()
        for shape in self.canvas.selectedShapes:
            self.addLabel(shape)
1863
        self.setDirty()
1864 1865 1866

    def moveShape(self):
        self.canvas.endMove(copy=False)
1867
        self.setDirty()
1868

D
Douglas Long 已提交
1869 1870 1871 1872
    def openDirDialog(self, _value=False, dirpath=None):
        if not self.mayContinue():
            return

K
Kentaro Wada 已提交
1873
        defaultOpenDirPath = dirpath if dirpath else "."
K
Kentaro Wada 已提交
1874
        if self.lastOpenDir and osp.exists(self.lastOpenDir):
D
Douglas Long 已提交
1875 1876
            defaultOpenDirPath = self.lastOpenDir
        else:
K
Kentaro Wada 已提交
1877 1878 1879
            defaultOpenDirPath = (
                osp.dirname(self.filename) if self.filename else "."
            )
D
Douglas Long 已提交
1880

K
Kentaro Wada 已提交
1881 1882 1883 1884 1885 1886 1887 1888 1889
        targetDirPath = str(
            QtWidgets.QFileDialog.getExistingDirectory(
                self,
                self.tr("%s - Open Directory") % __appname__,
                defaultOpenDirPath,
                QtWidgets.QFileDialog.ShowDirsOnly
                | QtWidgets.QFileDialog.DontResolveSymlinks,
            )
        )
D
Douglas Long 已提交
1890 1891
        self.importDirImages(targetDirPath)

1892 1893 1894 1895 1896 1897 1898 1899
    @property
    def imageList(self):
        lst = []
        for i in range(self.fileListWidget.count()):
            item = self.fileListWidget.item(i)
            lst.append(item.text())
        return lst

A
Aleksi J 已提交
1900 1901 1902
    def importDroppedImageFiles(self, imageFiles):
        self.filename = None
        for file in imageFiles:
K
Kentaro Wada 已提交
1903 1904 1905
            if file in self.imageList or not file.lower().endswith(
                tuple(EXTENSIONS)
            ):
A
Aleksi J 已提交
1906
                continue
K
Kentaro Wada 已提交
1907
            label_file = osp.splitext(file)[0] + ".json"
A
Aleksi J 已提交
1908 1909 1910 1911 1912
            if self.output_dir:
                label_file_without_path = osp.basename(label_file)
                label_file = osp.join(self.output_dir, label_file_without_path)
            item = QtWidgets.QListWidgetItem(file)
            item.setFlags(Qt.ItemIsEnabled | Qt.ItemIsSelectable)
K
Kentaro Wada 已提交
1913 1914 1915
            if QtCore.QFile.exists(label_file) and LabelFile.is_label_file(
                label_file
            ):
A
Aleksi J 已提交
1916 1917 1918 1919 1920 1921 1922 1923 1924 1925 1926
                item.setCheckState(Qt.Checked)
            else:
                item.setCheckState(Qt.Unchecked)
            self.fileListWidget.addItem(item)

        if len(self.imageList) > 1:
            self.actions.openNextImg.setEnabled(True)
            self.actions.openPrevImg.setEnabled(True)

        self.openNextImg()

K
Kentaro Wada 已提交
1927
    def importDirImages(self, dirpath, pattern=None, load=True):
1928 1929 1930
        self.actions.openNextImg.setEnabled(True)
        self.actions.openPrevImg.setEnabled(True)

D
Douglas Long 已提交
1931 1932 1933 1934 1935 1936
        if not self.mayContinue() or not dirpath:
            return

        self.lastOpenDir = dirpath
        self.filename = None
        self.fileListWidget.clear()
1937
        for filename in self.scanAllImages(dirpath):
K
Kentaro Wada 已提交
1938 1939
            if pattern and pattern not in filename:
                continue
K
Kentaro Wada 已提交
1940
            label_file = osp.splitext(filename)[0] + ".json"
1941
            if self.output_dir:
1942 1943
                label_file_without_path = osp.basename(label_file)
                label_file = osp.join(self.output_dir, label_file_without_path)
1944 1945
            item = QtWidgets.QListWidgetItem(filename)
            item.setFlags(Qt.ItemIsEnabled | Qt.ItemIsSelectable)
K
Kentaro Wada 已提交
1946 1947 1948
            if QtCore.QFile.exists(label_file) and LabelFile.is_label_file(
                label_file
            ):
1949 1950 1951
                item.setCheckState(Qt.Checked)
            else:
                item.setCheckState(Qt.Unchecked)
D
Douglas Long 已提交
1952
            self.fileListWidget.addItem(item)
1953
        self.openNextImg(load=load)
D
Douglas Long 已提交
1954 1955 1956 1957 1958

    def scanAllImages(self, folderPath):
        images = []
        for root, dirs, files in os.walk(folderPath):
            for file in files:
A
Aleksi J 已提交
1959
                if file.lower().endswith(tuple(EXTENSIONS)):
K
Kentaro Wada 已提交
1960
                    relativePath = osp.join(root, file)
1961
                    images.append(relativePath)
D
Douglas Long 已提交
1962 1963
        images.sort(key=lambda x: x.lower())
        return images