app.py

import io
import requests
import gradio as gr
# from transformers import AutoModel, AutoTokenizer
from huggingface_hub import list_models
from datasets import load_dataset
from typing import List
from PIL import Image


def get_image_names(dataset):
    return [str(i) for i in range(len(dataset))]

def get_image_from_dataset(index):
    image_data = dataset[int(index)]["image"]
    return image_data

def process_image(image=None, dataset_image_index=None):
    if dataset_image_index:
        image = get_image_from_dataset(dataset_image_index)

    return image



def create_interface(tag, image_indices):
    """ Create Gradio interface"""
    iface = gr.Interface(
        fn=process_image,
        inputs=[
            gr.Dropdown(choices=get_collection_models(tag), label="Select Model"),
            gr.Image(type="pil", label="Upload Image"),
            gr.Dropdown(choices=image_indices, label="Select one from MERIT Dataset test-set"),
        ],
        outputs=gr.Image(label="Output Image"),
        title="Saliency Visualization",
        description="Upload your image or select one from the MERIT Dataset test-set."
    )
    return iface


def get_collection_models(tag: str) -> List[str]:
    """Get a list of models from a specific Hugging Face collection."""
    models = list_models(author="de-Rodrigo")

    model_names = []
    for model in models:
        if tag in model.tags:
            model_names.append(model.modelId)

    return model_names

def load_model(model_name: str):
    """Load a model from Hugging Face Hub."""
    model = AutoModel.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    return model, tokenizer

# # Example processing function
# def process_input(text: str, model_name: str) -> str:
#     model, tokenizer = load_model(model_name)
#     inputs = tokenizer(text, return_tensors="pt")
#     outputs = model(**inputs)
#     return f"Processed output with {model_name}"


dataset_name = "de-Rodrigo/merit"
dataset = load_dataset(dataset_name, name="en-digital-seq", split="train", num_proc=8)
image_indices = get_image_names(dataset)

models_tag = "saliency-merit"

iface = create_interface(models_tag, image_indices)
iface.launch()