from imwatermark import WatermarkDecoder
import gradio as gr

# https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion_xl/watermark.py
# Copied from https://github.com/Stability-AI/generative-models/blob/613af104c6b85184091d42d374fef420eddb356d/scripts/demo/streamlit_helpers.py#L66
WATERMARK_MESSAGE = 0b101100111110110010010000011110111011000110011110
# bin(x)[2:] gives bits of x as str, use int to convert them to 0/1
WATERMARK_BITS = [int(bit) for bit in bin(WATERMARK_MESSAGE)[2:]]

SDXL_MESSAGE=[bool(bit) for bit in bin(WATERMARK_MESSAGE)[2:]]

def detect_watermark(image):
    decoder = WatermarkDecoder('bits', 48)
    watermark = decoder.decode(image, 'dwtDct')
    
    if(all(watermark==SDXL_MESSAGE)):
        return "この画像はdiffusersのSDXLで作られた可能性があります。"
    else:
        return "何もわかりませんでした。"

gr.Interface(fn=detect_watermark,
             title="AI画像簡易チェックツール",
             description="指定の画像が画像生成AIによって作られたものか電子透かしを読み取って簡易的に判定します。このツールで判定したことを決して断定的に用いないでください。[DALLE3とFirefly用](https://contentcredentials.org/verify)",
             inputs=gr.Image(type="numpy"),
             outputs=gr.Textbox()).launch()