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


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("----")
    '''
    return "done"

def classify(image, classifier_pipeline):
    return "done"

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
    '''
    chosen_model_name = st.selectbox("Select the model to use",  MODELS, index=0)
    if chosen_model_name is not None:
        st.write("You selected", chosen_model_name) 
        
    '''
        Dataset
    '''          
    shosen_dataset_name = st.selectbox("Select the dataset to use",  DATASETS, index=0)
    if shosen_dataset_name is not None:
        st.write("You selected", shosen_dataset_name)
        

           
    '''
        click to classify
        image_object = dataset['pasta'][0]
    '''    
    if chosen_model_name is not None and shosen_dataset_name is not None:
        if st.button("Classify images"):
            
            classification_obj1 =[]
            st.write("# FLAG 1")
            
            dataset = load_dataset(shosen_dataset_name,"testedata_readme")
            st.write("# FLAG 2")
            
            #Igame 
            image_object = dataset['pasta'][0]["image"]
            st.image(image_object, caption="Uploaded Image", use_column_width=True)
            st.write("# FLAG 3")
            
            #modle instance
            classifier_pipeline = pipeline('image-classification', model=chosen_model_name, device=0)
            #classifier_pipeline = pipeline('image-classification', model="nateraw/vit-age-classifier", device=0)
            st.write("# FLAG 4")

            #classification
            classification_result = classify(image_object, classifier_pipeline)
            st.write(classification_result)
            st.write("# FLAG 5")
            classification_obj1.append(classification_result)
            st.write("# FLAG 6")
            st.write(classification_obj1)
    

if __name__ == "__main__":
    main()