import gradio as gr
import gradio.components as grc
from wmdetection.models import get_watermarks_detection_model
from wmdetection.pipelines.predictor import WatermarksPredictor
import os, glob


model, transforms = get_watermarks_detection_model(
    'convnext-wm_1102_v2', 
    fp16=False, 
    cache_dir='model_files'
)
predictor = WatermarksPredictor(model, transforms, 'cpu')


def predict(image, threshold=0.5):
    result = predictor.predict_image_confidence(image)
    values = result.tolist()
    wm_flag = 1 if values[1] >= threshold else 0
    return 'watermarked' if wm_flag else 'clean', "%.4f" % values[1] # prints "watermarked"


examples = glob.glob(os.path.join('images', 'clean', '*'))
examples.extend(glob.glob(os.path.join('images', 'watermark', '*')))
examples = [[e, 0.5] for e in examples]
iface = gr.Interface(fn=predict, inputs=[grc.Image(type="pil"), grc.Number(label="threshold", default=0.5)], 
             examples=examples, outputs=[grc.Textbox(label="class"), grc.Textbox(label="wm_confidence")])
iface.launch()