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
            "Nunt/backup_leonardo_2024-02-06"
]
DATASETS = [
            "Nunt/testedata",
            "Nunt/backup_leonardo"       
]           
MAX_N_LABELS = 5
SPLIT_TO_CLASSIFY = 'pasta'

COLS = st.columns([0.75, 0.25]) 
SCROLLABLE_TEXT = COLS[1].container(height=500)


def classify_full_dataset(shosen_dataset_name, chosen_model_name):
    image_count = 0
     
    #modle instance
    classifier_pipeline = pipeline('image-classification', model=chosen_model_name)    
        
    #dataset
    dataset = load_dataset(shosen_dataset_name,"testedata_readme") 
     
    for i in range(len(dataset)):
        SCROLLABLE_TEXT.write("i-1:" + str(i-1))
        image_object = dataset['pasta'][i-1]["image"]
        SCROLLABLE_TEXT.image(image_object, caption="Uploaded Image", width=300)
        #classification
        classification_result = classifier_pipeline(image_object)
        SCROLLABLE_TEXT.write(classification_result)
        #TODO save classification result in dataset
        image_count += 1
        SCROLLABLE_TEXT.write(f"Image count" + str(image_count))
        #SCROLLABLE_TEXT.write(image_count)    
 


def main():
     
    COLS[0].write("# Bulk Image Classification App")
      
    #with CONTAINER_BODY:
    with COLS[0]:
        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 = COLS[0].selectbox("Select the model to use",  MODELS, index=0)
    if chosen_model_name is not None:
        COLS[0].write("You selected")
        COLS[0].write(chosen_model_name)
        
    #Dataset
    shosen_dataset_name = COLS[0].selectbox("Select the dataset to use",  DATASETS, index=0)
    if shosen_dataset_name is not None:
        COLS[0].write("You selected")
        COLS[0].write(shosen_dataset_name)
        
    #click to classify
    if chosen_model_name is not None and shosen_dataset_name is not None:
        if COLS[0].button("Classify images"):
        
            classify_full_dataset(shosen_dataset_name, chosen_model_name)
            COLS[0].write("Classification result {classification_result}")
            COLS[0].write("--- END ---")
            #COLS[0].write(classification_result)
 
                  
if __name__ == "__main__":
    main()