Spaces:
Running
Running
| 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 | |
| SPLIT_TO_CLASSIFY = 'pasta' | |
| COL1, COL2 = st.columns([3, 1]) | |
| CONTAINER_TOP = st.container() | |
| CONTAINER_BODY = st.container() | |
| #(image_object, classifier_pipeline) | |
| #def classify_one_image(classifier_model, dataset_to_classify): | |
| #classify_one_image(image_object, classifier_pipeline) | |
| def classify_one_image(classifier_model, dataset_to_classify): | |
| #image_object = dataset[SPLIT_TO_CLASSIFY][i]["image"] | |
| #st.image(image_object, caption="Uploaded Image", width=300) | |
| #for i in range(len(dataset_to_classify)): | |
| #for image in dataset_to_classify: | |
| #image_object = dataset[SPLIT_TO_CLASSIFY][i]["image"] | |
| #st.image(image_object, caption="Uploaded Image", width=300) | |
| #st.write(f"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") | |
| with COL2: | |
| #Image teste load | |
| image_object = dataset['pasta'][0]["image"] | |
| st.image(image_object, caption="Uploaded Image", width=300) | |
| st.write("### FLAG 3") | |
| #modle instance | |
| classifier_pipeline = pipeline('image-classification', model=chosen_model_name) | |
| st.write("### FLAG 4") | |
| #classification | |
| classification_result = classifier_pipeline(image_object) | |
| st.write(classification_result) | |
| st.write("### FLAG 5") | |
| #classification_array.append(classification_result) | |
| #save classification | |
| image_count += 1 | |
| return image_count | |
| def main(): | |
| with CONTAINER_TOP: | |
| st.title("Bulk Image Classification DEMO") | |
| # Restart or reset your app | |
| # if st.button("Restart"): | |
| # # Code to restart or reset your app goes here | |
| # import subprocess | |
| # subprocess.call(["shutdown", "-r", "-t", "0"]) | |
| CONTAINER_BODY = st.container(): | |
| with COL1: | |
| 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() |