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

#(image_object, classifier_pipeline)
#def classify_one_image(classifier_model, dataset_to_classify):
def classify_one_image(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_full_dataset(shosen_dataset_name, chosen_model_name):
    image_count = 0
     
    #dataset
    dataset = load_dataset(shosen_dataset_name,"testedata_readme")
    #Image teste load 
    image_object = dataset['pasta'][0]["image"]
    st.image(image_object, caption="Uploaded Image", height=300)
    st.write("### FLAG 3")
    
    #modle instance
    classifier_pipeline = pipeline('image-classification', model=chosen_model_name)
    st.write("### FLAG 4")
    
    #classification
    classification_result = classify_one_image(image_object, classifier_pipeline)
    st.write(classification_result)
    st.write("### FLAG 5")
    #classification_array.append(classification_result)
    
    #save classification
    
    
    image_count += 1
    
    return image_count




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_array =[]
            classification_result = classify_full_dataset(shosen_dataset_name, chosen_model_name)
            st.write(f"Classification result: {classification_result}")
            #classification_array.append(classification_result)
            #st.write("# FLAG 6")
            #st.write(classification_array)   

if __name__ == "__main__":
    main()