import gradio as gr

from transformers import pipeline, AutoModelForImageClassification, AutoFeatureExtractor

HF_MODEL_PATH = (
    "ImageIN/convnext-base-224_finetuned_on_unlabelled_IA_with_snorkel_labels"
)

classif_model = AutoModelForImageClassification.from_pretrained(HF_MODEL_PATH)
feature_extractor = AutoFeatureExtractor.from_pretrained(HF_MODEL_PATH)

classif_pipeline = pipeline(
    "image-classification", model=classif_model, feature_extractor=feature_extractor
)

OUTPUT_SENTENCE = "This image is {result}."


def get_formatted_prediction(img) -> str:
    return OUTPUT_SENTENCE.format(
        result=classif_pipeline(img)[0]["label"].replace("-", " ")
    )


demo = gr.Interface(
    fn=get_formatted_prediction,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="ImageIN",
    description="Identify illustrations in pages of historical books!",
    examples=["old_book_page.png", "women_book_image.png", "page_with_images.png"],
)
demo.launch()