|
import streamlit as st |
|
from transformers import pipeline |
|
from PIL import Image |
|
from datasets import load_dataset, Image |
|
from PIL import Image, ExifTags |
|
|
|
|
|
MODELS = [ |
|
"google/vit-base-patch16-224", |
|
"nateraw/vit-age-classifier" |
|
|
|
] |
|
DATASETS = [ |
|
"Nunt/testedata", |
|
"Nunt/backup_leonardo", |
|
"Nunt/backup_leonardo_2024-02-06" |
|
] |
|
MAX_N_LABELS = 5 |
|
SPLIT_TO_CLASSIFY = 'pasta' |
|
|
|
COLS = st.columns([0.60, 0.40]) |
|
SCROLLABLE_TEXT = COLS[1].container(height=500) |
|
|
|
|
|
def extract_file_name(image_object): |
|
file_name = image_object.filename |
|
return file_name |
|
|
|
|
|
def extract_index(file_name): |
|
return "flag: todo" |
|
|
|
|
|
def classify_full_dataset(shosen_dataset_name, chosen_model_name): |
|
image_count = 0 |
|
|
|
|
|
classifier_pipeline = pipeline('image-classification', model=chosen_model_name) |
|
|
|
|
|
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) |
|
|
|
|
|
file_name = extract_file_name(image_object) |
|
SCROLLABLE_TEXT.write(f"file_name: {file_name}") |
|
image_index = extract_index(file_name) |
|
SCROLLABLE_TEXT.write(f"image_index: {image_index}") |
|
|
|
|
|
|
|
image_object_texte_data = dataset['pasta'][i-1] |
|
SCROLLABLE_TEXT.write(f"image_object_texte_data: {image_object_texte_data}") |
|
|
|
|
|
|
|
SCROLLABLE_TEXT.write(f"teste1 (image_object._getexif()): {teste1}") |
|
|
|
teste2 = image_object._getexif().items() |
|
SCROLLABLE_TEXT.write(f"teste2 (image_object._getexif().items()): {teste2}") |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
classification_result = classifier_pipeline(image_object) |
|
SCROLLABLE_TEXT.write(classification_result) |
|
|
|
image_count += 1 |
|
SCROLLABLE_TEXT.write(f"Image count" + str(image_count)) |
|
|
|
|
|
|
|
|
|
def main(): |
|
|
|
COLS[0].write("# Bulk Image Classification App") |
|
|
|
|
|
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") |
|
|
|
|
|
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) |
|
|
|
|
|
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) |
|
|
|
|
|
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 ---") |
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
main() |