import cv2
from doc_ufcn.main import DocUFCN
import gradio as gr
from doc_ufcn import models
from PIL import Image
from PIL import ImageDraw

model_path, parameters = models.download_model("generic-page")

model = DocUFCN(len(parameters["classes"]), parameters["input_size"], "cpu")
model.load(model_path, parameters["mean"], parameters["std"])


def query_image(image):
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    detected_polygons, probabilities, mask, overlap = model.predict(
        image, raw_output=True, mask_output=True, overlap_output=True
    )
    max_confidence = 0
    max_prediction = None
    for prediction in detected_polygons[1]:
        if prediction["confidence"] > max_confidence:
            max_confidence = prediction["confidence"]
            max_prediction = prediction
    image = Image.fromarray(image)
    img2 = image.copy()
    xy = max_prediction["polygon"]
    draw = ImageDraw.Draw(img2)
    draw.polygon(xy, fill="blue")
    return Image.blend(image, img2, 0.5)


demo = gr.Interface(
    query_image,
    inputs=[gr.Image()],
    outputs="image",
    title="doc-ufcn Page Detection Demo",
    description="A demo showing the top prediction from the [Teklia/doc-ufcn-generic-page]("
    "https://huggingface.co/Teklia/doc-ufcn-generic-page) model. This generic page detection model predicts single pages from document images.",
    examples=[
        "v2_bsb00046516_00100_full_full_0_default.jpg",
        "512px-Page_from_the_Arthurian_Romances_illuminated_manuscript.jpg",
    ],
)

demo.launch()