feat: add exif data processing (saving and loading)

This commit is contained in:
Manuel Schmid 2024-02-04 23:29:48 +01:00
parent ceefba9b69
commit b281375ce2
No known key found for this signature in database
GPG Key ID: 32C4F7569B40B84B
3 changed files with 38 additions and 18 deletions

View File

@ -7,6 +7,7 @@ from pathlib import Path
import gradio as gr
from PIL import Image
import fooocus_version
import modules.config
import modules.sdxl_styles
from modules.flags import MetadataScheme, Performance, Steps
@ -518,16 +519,28 @@ def get_metadata_parser(metadata_scheme: MetadataScheme) -> MetadataParser:
raise NotImplementedError
def read_info_from_image(filepath) -> tuple[str | None, dict, MetadataScheme | None]:
def read_info_from_image(filepath) -> tuple[str | None, MetadataScheme | None]:
with Image.open(filepath) as image:
items = (image.info or {}).copy()
parameters = items.pop('parameters', None)
metadata_scheme = items.pop('fooocus_scheme', None)
exif = items.pop('exif', None)
if parameters is not None and is_json(parameters):
parameters = json.loads(parameters)
elif exif is not None:
exif = image.getexif()
# 0x9286 = UserComment
parameters = exif.get(0x9286, None)
# 0x927C = MakerNote
metadata_scheme = exif.get(0x927C, None)
if is_json(parameters):
parameters = json.loads(parameters)
try:
metadata_scheme = MetadataScheme(items.pop('fooocus_scheme', None))
metadata_scheme = MetadataScheme(metadata_scheme)
except ValueError:
metadata_scheme = None
@ -538,4 +551,16 @@ def read_info_from_image(filepath) -> tuple[str | None, dict, MetadataScheme | N
if isinstance(parameters, str):
metadata_scheme = MetadataScheme.A1111
return parameters, items, metadata_scheme
return parameters, metadata_scheme
def get_exif(metadata: str | None, metadata_scheme: str):
exif = Image.Exif()
# tags see see https://github.com/python-pillow/Pillow/blob/9.2.x/src/PIL/ExifTags.py
# 0x9286 = UserComment
exif[0x9286] = metadata
# 0x0131 = Software
exif[0x0131] = 'Fooocus v' + fooocus_version.version
# 0x927C = MakerNote
exif[0x927C] = metadata_scheme
return exif

View File

@ -6,9 +6,8 @@ import urllib.parse
from PIL import Image
from PIL.PngImagePlugin import PngInfo
from fooocus_version import version
from modules.util import generate_temp_filename
from modules.meta_parser import MetadataParser
from modules.meta_parser import MetadataParser, get_exif
log_cache = {}
@ -27,7 +26,8 @@ def log(img, metadata, metadata_parser: MetadataParser | None = None, output_for
date_string, local_temp_filename, only_name = generate_temp_filename(folder=path_outputs, extension=output_format)
os.makedirs(os.path.dirname(local_temp_filename), exist_ok=True)
parsed_parameters = metadata_parser.parse_string(metadata) if metadata_parser else None
parsed_parameters = metadata_parser.parse_string(metadata) if metadata_parser is not None else ''
image = Image.fromarray(img)
if output_format == 'png':
if parsed_parameters != '':
@ -36,14 +36,13 @@ def log(img, metadata, metadata_parser: MetadataParser | None = None, output_for
pnginfo.add_text('fooocus_scheme', metadata_parser.get_scheme().value)
else:
pnginfo = None
Image.fromarray(img).save(local_temp_filename, pnginfo=pnginfo)
image.save(local_temp_filename, pnginfo=pnginfo)
elif output_format == 'jpg':
Image.fromarray(img).save(local_temp_filename, quality=95, optimize=True, progressive=True)
image.save(local_temp_filename, quality=95, optimize=True, progressive=True, exif=get_exif(parsed_parameters, metadata_parser.get_scheme().value) if metadata_parser else Image.Exif())
elif output_format == 'webp':
Image.fromarray(img).save(local_temp_filename, quality=95, lossless=False)
image.save(local_temp_filename, quality=95, lossless=False, exif=get_exif(parsed_parameters, metadata_parser.get_scheme().value) if metadata_parser else Image.Exif())
else:
Image.fromarray(img).save(local_temp_filename)
image.save(local_temp_filename)
if args_manager.args.disable_image_log:
return local_temp_filename

View File

@ -278,15 +278,12 @@ with shared.gradio_root:
metadata_import_button = gr.Button(value='Apply Metadata')
def trigger_metadata_preview(filepath):
parameters, items, metadata_scheme = modules.meta_parser.read_info_from_image(filepath)
parameters, metadata_scheme = modules.meta_parser.read_info_from_image(filepath)
results = {}
if parameters is not None:
results['parameters'] = parameters
if items:
results['items'] = items
if isinstance(metadata_scheme, flags.MetadataScheme):
results['metadata_scheme'] = metadata_scheme.value
@ -326,8 +323,7 @@ with shared.gradio_root:
output_format = gr.Radio(label='Output Format',
choices=modules.flags.output_formats,
value=modules.config.default_output_format,
info='Metadata support has only been implemented for png.')
value=modules.config.default_output_format)
negative_prompt = gr.Textbox(label='Negative Prompt', show_label=True, placeholder="Type prompt here.",
info='Describing what you do not want to see.', lines=2,
@ -743,7 +739,7 @@ with shared.gradio_root:
load_parameter_button.click(modules.meta_parser.load_parameter_button_click, inputs=[prompt, state_is_generating], outputs=load_data_outputs, queue=False, show_progress=False)
def trigger_metadata_import(filepath, state_is_generating):
parameters, items, metadata_scheme = modules.meta_parser.read_info_from_image(filepath)
parameters, metadata_scheme = modules.meta_parser.read_info_from_image(filepath)
if parameters is None:
print('Could not find metadata in the image!')
parsed_parameters = {}