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


model, transforms = get_watermarks_detection_model(
    'convnext-tiny', 
    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=[gr.inputs.Image(type="pil"), gr.inputs.Number(label="threshold", default=0.5), ], 
             examples=examples, outputs=[gr.outputs.Textbox(label="class"), gr.outputs.Textbox(label="wm_confidence")])
iface.launch()