import gradio
from towhee.types.image_utils import from_pil
import towhee
import pandas as pd
from pymilvus import connections
import os

connections.connect(
    alias="default",
    host=os.getenv("milvus.host"),
    port=os.getenv("milvus.port"),
    user=os.getenv("milvus.user"),
    password=os.getenv("milvus.password")
)


with towhee.api() as api:
    image_search_function = (
        api.runas_op(func=lambda img: from_pil(img))
        .image_embedding.timm(model_name='resnet50')
        .tensor_normalize()
        .milvus_search(collection='reverse_image_search', limit=3)
        .runas_op(func=lambda res: [x.id for x in res])
        .as_function()
    )

interface = gradio.Interface(image_search_function,
                             gradio.inputs.Image(type="pil", source='upload'),
                             [gradio.outputs.Image(type="filepath", label=None) for _ in range(3)]
                             )
interface.launch(inline=True)