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 | |
| "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() |