import streamlit as st
from transformers import pipeline
from PIL import Image

MODEL_1 = "google/vit-base-patch16-224"
MIN_ACEPTABLE_SCORE = 0.1
MAX_N_LABELS = 5
MODEL_2 = "nateraw/vit-age-classifier"
MODELS = [
            "google/vit-base-patch16-224", #Classifição geral
            "nateraw/vit-age-classifier", #Classifição de idade
            "microsoft/resnet-50", #Classifição geral
            "Falconsai/nsfw_image_detection", #Classifição NSFW
            "cafeai/cafe_aesthetic", #Classifição de estética
            "microsoft/resnet-18", #Classifição geral
            "microsoft/resnet-34", #Classifição geral escolhida pelo copilot 
            "microsoft/resnet-101", #Classifição geral escolhida pelo copilot 
            "microsoft/resnet-152", #Classifição geral escolhida pelo copilot
            "microsoft/swin-tiny-patch4-window7-224",#Classifição geral
            "-- Reinstated on testing--",
            "microsoft/beit-base-patch16-224-pt22k-ft22k", #Classifição geral
            "-- New --"
            "-- Still in the testing process --"
            "facebook/convnext-large-224"
            "timm/resnet50.a1_in1k"
            "timm/mobilenetv3_large_100.ra_in1k"
            "trpakov/vit-face-expression"
            "rizvandwiki/gender-classification"
            "#q-future/one-align"
            "LukeJacob2023/nsfw-image-detector"
            "vit-base-patch16-224-in21k"
            "not-lain/deepfake"
            "carbon225/vit-base-patch16-224-hentai"
            "facebook/convnext-base-224-22k-1k"
            "facebook/convnext-large-224"
            "facebook/convnext-tiny-224"
            "nvidia/mit-b0"
            "microsoft/resnet-18"
            "microsoft/swinv2-base-patch4-window16-256"
            "andupets/real-estate-image-classification"
            "timm/tf_efficientnetv2_s.in21k"
            "timm/convnext_tiny.fb_in22k"
            "DunnBC22/vit-base-patch16-224-in21k_Human_Activity_Recognition"
            "FatihC/swin-tiny-patch4-window7-224-finetuned-eurosat-watermark"
            "aalonso-developer/vit-base-patch16-224-in21k-clothing-classifier"
            "RickyIG/emotion_face_image_classification"
            "shadowlilac/aesthetic-shadow"

        ]

def classify(image, model):
    classifier = pipeline("image-classification", model=model)
    result= classifier(image)
    return result

def save_result(result):
    st.write("In the future, this function will save the result in a database.")

def print_result(result):

    comulative_discarded_score = 0
    for i in range(len(result)):
        if result[i]['score'] < MIN_ACEPTABLE_SCORE:
            comulative_discarded_score += result[i]['score']
        else:
            st.write(result[i]['label'])
            st.progress(result[i]['score'])
            st.write(result[i]['score'])

    st.write(f"comulative_discarded_score:")
    st.progress(comulative_discarded_score)
    st.write(comulative_discarded_score)
    


def main():
    st.title("Image Classification")
    st.write("This is a simple web app to test and compare different image classifier models using Hugging Face's image-classification pipeline.")
    st.write("From time to time more models will be added to the list. If you want to add a model, please open an issue on the GitHub repository.")
    
    st.write("The models available are:")    
    shosen_model = st.selectbox("Select the model to use",  MODELS)
    
    st.write("Upload an image and click on the 'Classify' button to classify the image.")
    input_image = st.file_uploader("Upload Image")
    
    if input_image is not None:
        image_to_classify = Image.open(input_image)
        st.image(image_to_classify, caption="Uploaded Image", use_column_width=True)
        if st.button("Classify"):
            image_to_classify = Image.open(input_image)
            classification_obj1 =[]
            avable_models = st.selectbox
            
            classification_result = classify(image_to_classify, shosen_model)
            classification_obj1.append(classification_result)
            print_result(classification_result)
            save_result(classification_result)


if __name__ == "__main__":
    main()