app.py

import gradio as gr
from PIL import Image, ImageFilter
from PIL.ExifTags import TAGS
import numpy as np
import re


def main():
    with gr.Blocks() as iface:
        gr.Markdown("# Image Metadata Extractor")
        gr.Markdown("Upload an image (PNG, JPG, JPEG) to extract its metadata and parse it for prompts.\n")
        
        with gr.Row():
            with gr.Column():
                image_input = gr.File(label="Upload Image", type="filepath")
                censor_button = gr.Button("👁️ Toggle Censor")
                image_output = gr.Image(label="Image Preview")
                original_metadata_output = gr.Textbox(label="Full Metadata", lines=20)

            with gr.Column():
                prompt_output = gr.Textbox(label="Prompt", lines=4, show_copy_button=True)
                negative_prompt_output = gr.Textbox(label="Negative Prompt", lines=4, show_copy_button=True)
                seed_output = gr.Textbox(label="Seed Number", lines=1, show_copy_button=True)
                adetailer_prompt_1_output = gr.Textbox(label="ADetailer Prompt 1", lines=3, show_copy_button=True)
                adetailer_prompt_2_output = gr.Textbox(label="ADetailer Prompt 2", lines=3, show_copy_button=True)
                adetailer_negative_prompt_output = gr.Textbox(label="ADetailer Negative Prompt", lines=4, show_copy_button=True)
                original_prompt_output = gr.Textbox(label="Original Prompt", lines=4, show_copy_button=True)
                hires_prompt_output = gr.Textbox(label="Hires Prompt", lines=4, show_copy_button=True)

        image_state = gr.State(value={"original": None, "censored": None, "current": None})

        image_input.change(
            fn=extract_metadata,
            inputs=image_input,
            outputs=[
                image_output,
                prompt_output,
                negative_prompt_output,
                seed_output,
                adetailer_prompt_1_output,
                adetailer_prompt_2_output,
                adetailer_negative_prompt_output,
                original_prompt_output,
                hires_prompt_output,
                original_metadata_output,
                image_state,
            ],
        )

        censor_button.click(
            fn=toggle_censor,
            inputs=image_state,
            outputs=image_output,
        )

    iface.launch(show_error=True)

def decode_ansi_unix_exif(value):
    if not isinstance(value, bytes):
        return str(value)
    
    if value.startswith(b'UNICODE'):
        try:
            data = value[8:]
            
            result = []
            i = 0
            while i < len(data):
                if 32 <= data[i] <= 126:
                    result.append(chr(data[i]))
                if i+1 < len(data) and data[i+1] == 0:
                    i += 2
                else:
                    i += 1
            
            texto = ''.join(result)
            return texto
        except Exception as e:
            print(f"Error decoding ANSI/UNIX: {e}")
    
    for encoding in ['cp1252', 'latin-1', 'ascii', 'utf-8', 'utf-16le', 'utf-16be']:
        try:
            result = value.decode(encoding, errors='ignore')
            if result.strip():
                return result
        except:
            continue
    
    return ''.join(chr(b) for b in value if 32 <= b <= 126)

def extract_metadata(image_file):
    if image_file is None:
        return (
            None,  # image_output
            "N/A",  # prompt
            "N/A",  # negative_prompt
            -1,     # seed_number
            "N/A",  # adetailer_prompt_1
            "N/A",  # adetailer_prompt_2
            "N/A",  # adetailer_negative_prompt
            "N/A",  # original_prompt
            "N/A",  # hires_prompt
            "",     # metadata_str
            {"original": None, "censored": None, "current": None},  # image_state
        )
    
    img = Image.open(image_file)
    metadata = {}
    parameters_text = ""

    if img.format == "PNG":
        metadata = img.info
        if "parameters" in metadata:
            parameters_text = metadata["parameters"]
    elif img.format in ["JPEG"]:
        exif_data = img._getexif()
        if exif_data:
            for tag, value in exif_data.items():
                tag_name = TAGS.get(tag, tag)
                
                if tag_name == "UserComment" and isinstance(value, bytes):
                    decoded_value = decode_ansi_unix_exif(value)
                    value = decoded_value
                    
                    parameters_text = decoded_value
                elif isinstance(value, bytes):
                    value = decode_ansi_unix_exif(value)
                    
                metadata[tag_name] = value

    censored_img = img.filter(ImageFilter.GaussianBlur(40))
    img_display = np.array(img)

    prompt = "N/A"
    negative_prompt = "N/A"
    adetailer_prompt_1 = "N/A"
    adetailer_prompt_2 = "N/A"
    adetailer_negative_prompt = "N/A"
    original_prompt = "N/A"
    hires_prompt = "N/A"
    seed_number = -1

    metadata_str = "\n".join([f"{key}: {value}" for key, value in metadata.items()])

    if parameters_text:
        print(f"DEBUG - Parameters found: {parameters_text[:100]}...")
        
        prompt_match = re.search(r"(.*?)(?:Negative prompt:|Steps:|$)", parameters_text, re.DOTALL)
        if prompt_match:
            prompt = prompt_match.group(1).strip()

        negative_prompt_match = re.search(r"Negative prompt:(.*?)(?:Steps:|$)", parameters_text, re.DOTALL)
        if negative_prompt_match:
            negative_prompt = negative_prompt_match.group(1).strip()

        adetailer_prompt_1_match = re.search(r'ADetailer prompt:\s*"(.*?)"', parameters_text)
        if adetailer_prompt_1_match:
            adetailer_prompt_1 = adetailer_prompt_1_match.group(1).strip()

        adetailer_prompt_2_match = re.search(r'ADetailer negative prompt 2nd:\s*"(.*?)"', parameters_text)
        if adetailer_prompt_2_match:
            adetailer_prompt_2 = adetailer_prompt_2_match.group(1).strip()

        adetailer_negative_prompt_match = re.search(r'ADetailer negative prompt:\s*"(.*?)"', parameters_text)
        if adetailer_negative_prompt_match:
            adetailer_negative_prompt = adetailer_negative_prompt_match.group(1).strip()

        original_prompt_match = re.search(r'DTG prompt:\s*"(.*?)"', parameters_text, re.DOTALL)
        if original_prompt_match:
            original_prompt = original_prompt_match.group(1).strip()

        hires_prompt_match = re.search(r'Hires prompt:\s*"(.*?)"', parameters_text, re.DOTALL)
        if hires_prompt_match:
            hires_prompt = hires_prompt_match.group(1).strip()

        seed_match = re.search(r"Seed:\s*(\d+)", parameters_text)
        if seed_match:
            seed_number = seed_match.group(1).strip()
        
        if prompt == "N/A" and parameters_text.strip():
            prompt = parameters_text.strip()

    return (
        img_display,
        prompt,
        negative_prompt,
        seed_number,
        adetailer_prompt_1,
        adetailer_prompt_2,
        adetailer_negative_prompt,
        original_prompt,
        hires_prompt,
        metadata_str,
        {"original": img, "censored": censored_img, "current": img},
    )


def toggle_censor(image_state):
    if image_state["current"] == image_state["original"]:
        image_state["current"] = image_state["censored"]
    else:
        image_state["current"] = image_state["original"]
    
    return np.array(image_state["current"])


if __name__ == "__main__":
    main()