Update app.py

app.py

@@ -1,9 +1,20 @@
-from imwatermark import WatermarkEncoder
+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:]]
+
 def detect_watermark(image):
+    decoder = WatermarkDecoder('bytes', 32)
+    watermark = decoder.decode(image, 'dwtDct')
+    print(watermark.decode())
     return "何もわかりませんでした。"
     #return "この画像はdiffusersで作られた可能性があります。"
 gr.Interface(fn=detect_watermark,
-             inputs=gr.Image(),
+             title="AI画像簡易チェックツール",
+             description="指定の画像が画像生成AIによって作られたものかを簡易的に判定します。このツールで判定したことを決して断定的に用いないでください。",
+             inputs=gr.Image(type="numpy"),
              outputs=gr.Textbox()).launch()