Tag-Machine/Model/ImageWriter.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
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
from iptcinfo3 import IPTCInfo
import os
import exiv2

class ImageWriter:

    _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)