Nuno-Tome's picture
1st version is ready to teste
cc65489
raw
history blame
No virus
2.1 kB
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/vit-base-patch16-224", #Classifição geral
"NunT/vit-age-classifier" #Classifição de idade
]
MAX_N_LABELS = 5
def classify_images(classifier_model, dataset_to_classify):
for image in dataset:
st("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("----")
'''
def main():
st.title("Bulk Image Classification")
st.markdown("This app uses several 🤗 models to classify images stored in 🤗 datasets.")
st.write("Soon we will have a dataset template")
st.write("**Soon we will have dataset selector**")
'''
Model
'''
shosen_model_name = st.selectbox("Select the model to use", MODELS)
if shosen_model is not None:
st.write("You selected", shosen_model_name)
'''
Dataset
'''
shosen_dataset_name =st.radio("Select the model to use", MODELS)
if shosen_dataset is not None:
st.write("You selected", shosen_dataset_name)
image_object = dataset['pasta'][0]["image"]
if shosen_model is not None and shosen_dataset is not None:
st.image(shosen_dataset[0], caption="Uploaded Image", use_column_width=True)
if st.button("Classify images"):
dataset = load_dataset("Nunt/testedata","testedata_readme")
classifier = pipeline('image-classification', model=model_name, device=0)
classify_images(classifier, dataset)
if __name__ == "__main__":
main()