Spaces:
Running
Running
File size: 2,752 Bytes
c2ddf09 14eecb6 c2ddf09 14eecb6 c2ddf09 14eecb6 c2ddf09 14eecb6 c2ddf09 14eecb6 |
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 |
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 = {}
# Handling multiple possible metadata keys
potential_keys = ['metadata', 'prompt', 'Comment', 'parameters', 'exif']
for key in potential_keys:
if key in image.info:
raw_data = image.info[key]
# If raw data starts with '{', assume JSON format
if raw_data.startswith('{'):
metadata = json.loads(raw_data)
else:
if key == 'parameters':
# Attempt to process Stable Diffusion or NovelAI style data
lines = raw_data.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 ':' in line:
key, value = line.split(':', 1)
metadata[key.strip()] = value.strip()
elif key == 'Comment':
# Specific handling for NovelAI
metadata = json.loads(raw_data)
metadata['model'] = 'NovelAI'
break # Exit loop once a supported key is found
if not metadata:
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")
# Event listener for when the image is changed
input_image.change(
fn=process_image,
inputs=input_image,
outputs=[status_output, output_metadata],
api_name="interrogate"
)
demo.launch()
|