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