File size: 3,145 Bytes
c2ddf09 c30b9c5 c2ddf09 |
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 |
import gradio as gr
from PIL import Image, PngImagePlugin
import json
import traceback
def extract_metadata(image):
if image is None:
return "Please upload an image.", {}
try:
metadata = {}
if 'metadata' in image.info:
metadata = json.loads(image.info['metadata'])
elif 'prompt' in image.info:
metadata = json.loads(image.info['prompt'])
elif 'Comment' in image.info:
metadata = json.loads(image.info['Comment'])
metadata['model'] = 'NovelAI'
elif 'parameters' in image.info:
if image.info['parameters'].startswith('{'):
parameters_data = json.loads(image.info['parameters'])
if 'sui_image_params' in parameters_data:
sui_image_params = parameters_data['sui_image_params']
metadata.update(sui_image_params)
else:
metadata = parameters_data
else:
lines = image.info['parameters'].split('\n')
prompt = lines[0].strip()
negative_prompt = lines[1].strip().replace('Negative prompt:', '').strip()
metadata['prompt'] = prompt
metadata['negative_prompt'] = negative_prompt
for line in lines[2:]:
line = line.strip()
if line.startswith('Steps:'):
steps_info = line.split(':', 1)[1].strip().split(',')
for info in steps_info:
info = info.strip()
if ':' in info:
key, value = info.split(':', 1)
metadata[key.strip()] = value.strip()
else:
return "No supported metadata found in the image.", {}
return "Metadata extracted successfully.", metadata
except Exception as e:
error_message = f"Error extracting metadata: {str(e)}\n{traceback.format_exc()}"
return error_message, {}
def process_image(image):
status, metadata = extract_metadata(image)
return status, metadata
with gr.Blocks() as demo:
gr.Markdown(
"""
# Image Metadata Extractor
Extract and display metadata from images generated by various AI tools.
"""
)
with gr.Row():
with gr.Column():
input_image = gr.Image(label="Input Image", type="pil", height=480)
with gr.Column():
status_output = gr.Textbox(label="Status")
output_metadata = gr.JSON(label="Metadata")
input_image.change(
fn=process_image,
inputs=input_image,
outputs=[status_output, output_metadata],
api_name="interrogate"
)
gr.Examples(
examples=[
["example/asukatest.png"],
["example/arimakana.png"],
["example/stelle.png"],
["example/shinji.png"],
],
inputs=input_image,
outputs=[status_output, output_metadata],
fn=process_image,
cache_examples=True,
)
demo.launch() |