import streamlit as st
from transformers import pipeline
from PIL import Image
from datasets import load_dataset, Image, list_datasets
from PIL import Image

MODELS = [
            "google/vit-base-patch16-224", #Classifição geral
            "nateraw/vit-age-classifier" #Classifição de idade
]
DATASETS = [
            "Nunt/testedata",
            "Nunt/backup_leonardo_2024-02-01" 
]
MAX_N_LABELS = 5


def classify_images(classifier_model, dataset_to_classify):
       
    for image in dataset:
        st("Image classification: ", image['file'])
    '''
        image_path = image['file']
        img = Image.open(image_path)
        st.image(img, caption="Original image", use_column_width=True)
        results = classifier(image_path, top_k=MAX_N_LABELS)
        st.write(results)
        st.write("----")
    '''



def main():
    st.title("Bulk Image Classification")
    st.markdown("This app uses several 🤗 models to classify images stored in 🤗 datasets.")
    st.write("Soon we will have a dataset template")
    

    '''
        Model
    '''
    shosen_model_name = st.selectbox("Select the model to use",  MODELS, index=0)
    if shosen_model_name is not None:
        st.write("You selected", shosen_model_name) 
        
    '''
        Dataset
    '''          
    shosen_dataset_name = st.radio("Select the model to use",  DATASETS index=0)
    if shosen_dataset_name is not None:
        st.write("You selected", shosen_dataset_name)
        
        #image_object = dataset['pasta'][0]
           
    '''
        click to classify
    '''    
    if shosen_model_name is not None and shosen_dataset_name is not None:
        if st.button("Classify images"):
            dataset = load_dataset(shosen_dataset_name,"testedata_readme")
            st.image(dataset[0], caption="Uploaded Image", use_column_width=True)
            classifier = pipeline('image-classification', model=shosen_model_name, device=0)
            classify_images(classifier, dataset)


    

if __name__ == "__main__":
    main()