diff --git a/Model/ImageWriter.py b/Model/ImageWriter.py index 9852523..4fdf7fc 100644 --- a/Model/ImageWriter.py +++ b/Model/ImageWriter.py @@ -1,18 +1,128 @@ from iptcinfo3 import IPTCInfo +import os +import exiv2 class ImageWriter: - def writeTagsFromPredictionsInImages(self, predictions): + + _max_bytes = { + 'Iptc.Application2.Byline' : 32, + 'Iptc.Application2.BylineTitle' : 32, + 'Iptc.Application2.Caption' : 2000, + 'Iptc.Application2.City' : 32, + 'Iptc.Application2.Contact' : 128, + 'Iptc.Application2.Copyright' : 128, + 'Iptc.Application2.CountryCode' : 3, + 'Iptc.Application2.CountryName' : 64, + 'Iptc.Application2.Credit' : 32, + 'Iptc.Application2.Headline' : 256, + 'Iptc.Application2.Keywords' : 64, + 'Iptc.Application2.ObjectName' : 64, + 'Iptc.Application2.Program' : 32, + 'Iptc.Application2.ProgramVersion' : 10, + 'Iptc.Application2.ProvinceState' : 32, + 'Iptc.Application2.SpecialInstructions': 256, + 'Iptc.Application2.SubLocation' : 32, + 'Iptc.Envelope.CharacterSet' : 32, + } + + + def changeToRawExtension(self, file_path, newExtension): + base_path, extension = os.path.splitext(file_path) + raw_file_path = base_path + "." + newExtension + return raw_file_path + + @classmethod + def truncate_iptc(cls, tag, value): + if tag in cls._max_bytes: + value = value.encode('utf-8')[:cls._max_bytes[tag]] + value = value.decode('utf-8', errors='ignore') + return value + + def set_iptc_value(self, iptcData, tag, value): + if not value: + self.clear_iptc_tag(tag) + return + # make list of values + key = exiv2.IptcKey(tag) + type_id = exiv2.IptcDataSets.dataSetType(key.tag(), key.record()) + if type_id == exiv2.TypeId.date: + values = [exiv2.DateValue(*value)] + elif type_id == exiv2.TypeId.time: + values = [exiv2.TimeValue(*value)] + elif isinstance(value, (list, tuple)): + values = [self.truncate_iptc(tag, x) for x in value] + else: + values = [self.truncate_iptc(tag, value)] + # update or delete existing values + datum = iptcData.findKey(key) + while datum != iptcData.end(): + if datum.key() == tag: + if values: + datum.setValue(values.pop(0)) + else: + datum = iptcData.erase(datum) + continue + next(datum) + # append remaining values + while values: + datum = exiv2.Iptcdatum(key) + datum.setValue(values.pop(0)) + if iptcData.add(datum) != 0: + print('duplicated tag %s', tag) + return + return iptcData + + def writeTagsFromPredictionsInImages(self, predictions, applyToRaw, overwrite): + startApplyToRaw = applyToRaw for data in predictions: + applyToRaw = startApplyToRaw filename = data[0] + filenameRaw = None prediction = data[1] iptcInfo = IPTCInfo(filename, force=True) + iptcInfoRaw = None + if applyToRaw: + extensions = ['dng','DNG'] + for extension in extensions: + filenameRaw = self.changeToRawExtension(filename, extension) + if os.path.exists(filenameRaw): + iptcInfoRaw = exiv2.ImageFactory.open(filenameRaw) + iptcInfoRaw.readMetadata() + break + applyToRaw = iptcInfoRaw != None + + addedKeywords = [] + + if overwrite: + iptcInfo['keywords'] = [] if prediction.m_label not in iptcInfo['keywords']: iptcInfo['keywords'].append(prediction.m_label) + addedKeywords.append(prediction.m_label) for category in prediction.m_categories: if category not in iptcInfo['keywords']: iptcInfo['keywords'].append(category) + addedKeywords.append(category) + + try: + iptcInfo.save() + tempFilename = filename + "~" + if os.path.exists(tempFilename): + os.remove(tempFilename) + except Exception as e: + print('Error in file "'+filename+'" : \n\t', e) + + if applyToRaw: + try: + separator = '; ' + iptcData = iptcInfoRaw.iptcData() + iptcData = self.set_iptc_value(iptcData, 'Iptc.Envelope.CharacterSet', '\x1b%G') + iptcData = self.set_iptc_value(iptcData,"Iptc.Application2.Keywords",addedKeywords) + iptcInfoRaw.setIptcData(iptcData) + iptcInfoRaw.writeMetadata() + except Exception as e: + print('Error in file "'+filenameRaw+'" : \n\t', e) - iptcInfo.save() \ No newline at end of file + \ No newline at end of file diff --git a/UI/MainWindow.py b/UI/MainWindow.py index 9471c6e..8c5b051 100644 --- a/UI/MainWindow.py +++ b/UI/MainWindow.py @@ -15,6 +15,8 @@ def __init__(self, parent=None): # 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...") @@ -31,8 +33,14 @@ def __init__(self, parent=None): 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_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) @@ -59,5 +67,7 @@ def onClassifyImageButtonClicked(self): @QtCore.Slot() def onWriteTagsButtonClicked(self): lastPredictions = self.m_predictionModel.getLastPredictions() - self.m_imageWriter.writeTagsFromPredictionsInImages(lastPredictions) + applyToRaw = self.m_writeTagsCheckBox.isChecked() + overwrite = self.m_overwriteTagsCheckBox.isChecked() + self.m_imageWriter.writeTagsFromPredictionsInImages(lastPredictions, applyToRaw, overwrite) \ No newline at end of file