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