Tag-Machine/UI/MainWindow.py at f7a7ec904b13ba0670a4804dcacb786d56f76c98 · torresflo/Tag-Machine · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
from PySide6 import QtCore, QtWidgets, QtGui
from UI.PredictionItemModel import PredictionItemModel

from Model.PredictionModel import PredictionModel
from Model.ImageWriter import ImageWriter

class MainWindow(QtWidgets.QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)

        #Model
        self.m_predictionModel = PredictionModel()
        self.m_imageWriter = ImageWriter()

        # UI
        self.m_loadFileButton = QtWidgets.QPushButton("Load files...")
        self.m_classifyImageButton = QtWidgets.QPushButton("Classify images")
        self.m_writeTagsCheckBox = QtWidgets.QCheckBox("Apply to DNG")
        self.m_overwriteTagsCheckBox = QtWidgets.QCheckBox("Overwrite tags")
        self.m_writeTagsButton = QtWidgets.QPushButton("Write tags in images")

        self.m_fileNamesPlainTextEdit = QtWidgets.QPlainTextEdit("No files loaded...")
        self.m_fileNamesPlainTextEdit.setReadOnly(True)
        self.m_fileNamesPlainTextEdit.setMaximumHeight(100)
        self.m_fileNamesPlainTextEdit.setVerticalScrollBarPolicy(QtCore.Qt.ScrollBarAsNeeded)

        self.m_resultTable = QtWidgets.QTableView()
        self.m_resultTable.horizontalHeader().setSectionResizeMode(QtWidgets.QHeaderView.ResizeToContents)
        self.m_resultTable.horizontalHeader().setStretchLastSection(True)
        self.m_resultTable.verticalHeader().setSectionResizeMode(QtWidgets.QHeaderView.ResizeToContents)
        self.m_resultTable.setSelectionMode(QtWidgets.QAbstractItemView.NoSelection)

        self.layout = QtWidgets.QVBoxLayout(self)
        self.layout.addWidget(self.m_loadFileButton)
        self.layout.addWidget(self.m_fileNamesPlainTextEdit)
        self.layout.addWidget(self.m_classifyImageButton)
        self.layout.addWidget(self.m_resultTable)

        buttonLayout = QtWidgets.QHBoxLayout()
        buttonLayout.addWidget(self.m_writeTagsCheckBox)
        buttonLayout.addWidget(self.m_overwriteTagsCheckBox)
        self.layout.addLayout(buttonLayout)

        self.layout.addWidget(self.m_writeTagsButton)

        self.m_loadFileButton.clicked.connect(self.onLoadFileButtonClicked)
        self.m_classifyImageButton.clicked.connect(self.onClassifyImageButtonClicked)
        self.m_writeTagsButton.clicked.connect(self.onWriteTagsButtonClicked)

    @QtCore.Slot()
    def onLoadFileButtonClicked(self):
        self.m_fileNames, selectedFilter = QtWidgets.QFileDialog.getOpenFileNames(self, "Please select images", "", "Images (*.png *.jpg)")
        fileNamesString = ""
        for fileName in self.m_fileNames:
            fileNamesString += QtCore.QFileInfo(fileName).fileName() + '\n'
        self.m_fileNamesPlainTextEdit.setPlainText(fileNamesString)

    @QtCore.Slot()
    def onClassifyImageButtonClicked(self):
        predictions = self.m_predictionModel.computePredictions(self.m_fileNames)

        if(len(predictions) > 0):
            self.m_resultTable.setModel(PredictionItemModel(predictions))
        else:
            self.m_resultTable.setModel(QtGui.QStandardItemModel(0, 0))

    @QtCore.Slot()
    def onWriteTagsButtonClicked(self):
        lastPredictions = self.m_predictionModel.getLastPredictions()
        applyToRaw = self.m_writeTagsCheckBox.isChecked()
        overwrite = self.m_overwriteTagsCheckBox.isChecked()
        self.m_imageWriter.writeTagsFromPredictionsInImages(lastPredictions, applyToRaw, overwrite)