import gradio as gr
import re
from pathlib import Path
from api import load_annoy_index, analyze_image

annoy_index, mappings = load_annoy_index()

def get_article_text():
    article = Path("README.md").read_text()
    # Remove the HuggingFace Space app information from the README
    article = re.sub(r"^---.+---\s+", "", article, flags=re.MULTILINE + re.DOTALL)
    return article



def find_matching_images(input_img, n_matches: int = 10):
    results = analyze_image(input_img, annoy_index, n_matches=n_matches)
    

    indexes = results[0]
    # scores = results[1]

    images = []

    for i in range(len(indexes)):
        index = str(indexes[i])

        mapping = mappings[index]

        url = mapping["url"]
        if url != "":
          images.append(url)

    return images

iface = gr.Interface(
    find_matching_images,
    title="類似貼り込み資料検索",
    description="""類似する貼り込み資料を検索します。 Upload a picture and find out!
    Give it a shot or try one of the sample images below.
    
    Built with ❤️ using great open-source libraries such as PyTorch and Annoy.""",
    article=get_article_text(),
    inputs=[
        gr.inputs.Image(shape=None, label="Your image"),
    ],
    outputs=gr.Gallery(label="類似する貼り込み資料"),
    examples=[
        ["images/025_0085.jpg"],
        ["images/046_0051.jpg"],
    ],
)

iface.launch()