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, 'cuda:0')


def predict(image):
    result = predictor.predict_image(image)
    return 'watermarked' if result else 'clean' # prints "watermarked"


examples = glob.glob(os.path.join('images', 'clean', '*'))
examples.extend(glob.glob(os.path.join('images', 'watermark', '*')))
iface = gr.Interface(fn=predict, inputs=[gr.inputs.Image(type="pil")], 
             examples=examples, outputs="text")
iface.launch()