| 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: | |
| 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() |