import gradio as gr

from demo import query_image

description = """
Github link: <a href="https://github.com/ngthanhtin/owlvit_segment_anything">Link</a>
Gradio demo for combining <a href="https://github.com/facebookresearch/segment-anything">Segment-Anything (SAM)</a>
<a href="https://huggingface.co/docs/transformers/main/en/model_doc/owlvit">OWL-ViT</a>.
\n\nYou can use OWL-ViT to query boxes with text descriptions of any object, then SAM will segment anything in the boxes.

Note: In this demo, I only get the one box per class to demonstrate, you can modify the code to get multiple boxes beyond a threshold.
"""
demo = gr.Interface(
    query_image, 
    inputs=[gr.Image(), "text"], 
    outputs=["image", "image"],
    title="Segment Anything (SAM) with OWL-ViT",
    description=description,
    examples=[
        ["./demo_images/cats.png", "cats,ears"], 
        ["./demo_images/demo1.jpg", "bear,soil,sea"],
        ["./demo_images/demo2.jpg", "dog,ear,leg,eyes,tail"],
        ["./demo_images/tanager.jpg", "wing,eyes,back,legs,tail"]
    ],
)

# demo.launch()
demo.launch(server_name="0.0.0.0", debug=True)