automatic/cli/image-exif.py

#!/bin/env python

import os
import io
import re
import sys
import json
from PIL import Image, ExifTags, TiffImagePlugin, PngImagePlugin
from rich import print # pylint: disable=redefined-builtin


def unquote(text):
    if len(text) == 0 or text[0] != '"' or text[-1] != '"':
        return text
    try:
        return json.loads(text)
    except Exception:
        return text


def parse_generation_parameters(infotext):
    if not isinstance(infotext, str):
        return {}
    re_param = re.compile(r'\s*([\w ]+):\s*("(?:\\"[^,]|\\"|\\|[^\"])+"|[^,]*)(?:,|$)') # multi-word: value
    re_size = re.compile(r"^(\d+)x(\d+)$") # int x int
    basic_params = ['steps', 'seed', 'width', 'height', 'sampler', 'size', 'cfg scale', 'hires'] # first param is one of those

    sanitized = infotext.replace('prompt:', 'Prompt:').replace('negative prompt:', 'Negative prompt:').replace('Negative Prompt', 'Negative prompt') # cleanup everything in brackets so re_params can work
    sanitized = re.sub(r'<[^>]*>', lambda match: ' ' * len(match.group()), sanitized)
    sanitized = re.sub(r'\([^)]*\)', lambda match: ' ' * len(match.group()), sanitized)
    sanitized = re.sub(r'\{[^}]*\}', lambda match: ' ' * len(match.group()), sanitized)

    params = dict(re_param.findall(sanitized))
    params = { k.strip():params[k].strip() for k in params if k.lower() not in ['hashes', 'lora', 'embeddings', 'prompt', 'negative prompt']} # remove some keys
    if len(list(params)) == 0:
        first_param = None
    else:
        try:
            first_param, first_param_idx = next((s, i) for i, s in enumerate(params) if any(x in s.lower() for x in basic_params))
        except Exception:
            first_param, first_param_idx = next(iter(params)), 0
        if first_param_idx > 0:
            for _i in range(first_param_idx):
                params.pop(next(iter(params)))
    params_idx = sanitized.find(f'{first_param}:') if first_param else -1
    negative_idx = infotext.find("Negative prompt:")

    prompt = infotext[:params_idx] if negative_idx == -1 else infotext[:negative_idx] # prompt can be with or without negative prompt
    negative = infotext[negative_idx:params_idx] if negative_idx >= 0 else ''

    for k, v in params.copy().items(): # avoid dict-has-changed
        if len(v) > 0 and v[0] == '"' and v[-1] == '"':
            v = unquote(v)
        m = re_size.match(v)
        if v.replace('.', '', 1).isdigit():
            params[k] = float(v) if '.' in v else int(v)
        elif v == "True":
            params[k] = True
        elif v == "False":
            params[k] = False
        elif m is not None:
            params[f"{k}-1"] = int(m.group(1))
            params[f"{k}-2"] = int(m.group(2))
        elif k == 'VAE' and v == 'TAESD':
            params["Full quality"] = False
        else:
            params[k] = v
    params["Prompt"] = prompt.replace('Prompt:', '').strip()
    params["Negative prompt"] = negative.replace('Negative prompt:', '').strip()
    return params


class Exif: # pylint: disable=single-string-used-for-slots
    __slots__ = ('__dict__') # pylint: disable=superfluous-parens
    def __init__(self, image = None):
        super(Exif, self).__setattr__('exif', Image.Exif()) # pylint: disable=super-with-arguments
        self.pnginfo = PngImagePlugin.PngInfo()
        self.tags = {**dict(ExifTags.TAGS.items()), **dict(ExifTags.GPSTAGS.items())}
        self.ids = {**{v: k for k, v in ExifTags.TAGS.items()}, **{v: k for k, v in ExifTags.GPSTAGS.items()}}
        if image is not None:
            self.load(image)

    def __getattr__(self, attr):
        if attr in self.__dict__:
            return self.__dict__[attr]
        return self.exif.get(attr, None)

    def load(self, img: Image):
        img.load() # exif may not be ready
        exif_dict = {}
        try:
            exif_dict = dict(img._getexif().items()) # pylint: disable=protected-access
        except Exception:
            pass
        if not exif_dict:
            exif_dict = dict(img.info.items())
        for key, val in exif_dict.items():
            if isinstance(val, bytes): # decode bytestring
                val = self.decode(val)
            if val is not None:
                if isinstance(key, str):
                    self.exif[key] = val
                    self.pnginfo.add_text(key, str(val), zip=False)
                elif isinstance(key, int) and key in ExifTags.TAGS: # add known tags
                    if self.tags[key] in ['ExifOffset']:
                        continue
                    self.exif[self.tags[key]] = val
                    self.pnginfo.add_text(self.tags[key], str(val), zip=False)
                    # if self.tags[key] == 'UserComment': # add geninfo from UserComment
                        # self.geninfo = val
                else:
                    print('metadata unknown tag:', key, val)
        for key, val in self.exif.items():
            if isinstance(val, bytes): # decode bytestring
                self.exif[key] = self.decode(val)

    def decode(self, s: bytes):
        remove_prefix = lambda text, prefix: text[len(prefix):] if text.startswith(prefix) else text # pylint: disable=unnecessary-lambda-assignment
        for encoding in ['utf-8', 'utf-16', 'ascii', 'latin_1', 'cp1252', 'cp437']: # try different encodings
            try:
                s = remove_prefix(s, b'UNICODE')
                s = remove_prefix(s, b'ASCII')
                s = remove_prefix(s, b'\x00')
                val = s.decode(encoding, errors="strict")
                val = re.sub(r'[\x00-\x09]', '', val).strip() # remove remaining special characters
                if len(val) == 0: # remove empty strings
                    val = None
                return val
            except Exception:
                pass
        return None

    def parse(self):
        x = self.exif.pop('parameters', None) or self.exif.pop('UserComment', None)
        res = parse_generation_parameters(x)
        return res

    def get_bytes(self):
        ifd = TiffImagePlugin.ImageFileDirectory_v2()
        exif_stream = io.BytesIO()
        for key, val in self.exif.items():
            if key in self.ids:
                ifd[self.ids[key]] = val
            else:
                print('metadata unknown exif tag:', key, val)
        ifd.save(exif_stream)
        raw = b'Exif\x00\x00' + exif_stream.getvalue()
        return raw


def read_exif(filename: str):
    if filename.lower().endswith('.heic'):
        from pi_heif import register_heif_opener
        register_heif_opener()
    try:
        image = Image.open(filename)
        exif = Exif(image)
        print('image:', filename, 'format:', image)
        print('exif:', vars(exif.exif)['_data'])
        print('info:', exif.parse())
    except Exception as e:
        print('metadata error reading:', filename, e)


if __name__ == '__main__':
    sys.argv.pop(0)
    if len(sys.argv) == 0:
        print('metadata:', 'no files specified')
    for fn in sys.argv:
        if os.path.isfile(fn):
            read_exif(fn)
        elif os.path.isdir(fn):
            for root, _dirs, files in os.walk(fn):
                for file in files:
                    read_exif(os.path.join(root, file))