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/vit-base-patch16-224", #Classifição geral
            "NunT/vit-age-classifier" #Classifição de idade
]
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")
    st.write("**Soon we will have dataset selector**")   
    

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


    

if __name__ == "__main__":
    main()