from ultralytics import YOLO
import PIL
import gradio as gr
import numpy as np
import os

model = YOLO("best.pt")

def predict(input_img) -> tuple[np.ndarray | PIL.Image.Image | str, list[tuple[np.ndarray | tuple[int, int, int, int], str]]]:
    res = model(input_img)
    if len(res) == 0:
        return input_img, "No watermark detected"
    res = res[0]
    # convert res.boxes.xyxy to a tuple of (x1, y1, x2, y2)
    bbox = res.boxes.xyxy[0].tolist()
    bbox = (int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3]))
    # convert res.boxes.cls to a string
    label = res.boxes.cls[0]
    str_label = "Watermark is a logo" if label == 0 else "Watermark is a text"
    print(bbox, str_label)
    return input_img, [(bbox, str_label)]


gradio_app = gr.Interface(
    predict,
    inputs=gr.Image(label="Upload your watermaked image", sources=['upload'], type="pil"),
    # output displays the image with the bounding boxes
    outputs=gr.AnnotatedImage(),
    title="Detect Watermark in Images",
    description="This demo use a YoloV8 Nano model from Ultralytics, fine-tuned on the PITA Dataset for watermarked images",
    examples=[
        os.path.join(os.path.dirname(__file__), "samples/example_text1.jpg"),
        os.path.join(os.path.dirname(__file__), "samples/example_text2.jpg"),
        os.path.join(os.path.dirname(__file__), "samples/example_text3.jpg"),
        os.path.join(os.path.dirname(__file__), "samples/example_logo1.jpg"),
        os.path.join(os.path.dirname(__file__), "samples/example_logo2.jpg"),
        os.path.join(os.path.dirname(__file__), "samples/example_logo3.jpg"),
    ],
    allow_flagging="never"
)


if __name__ == "__main__":
    gradio_app.launch()