|
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", |
|
"nateraw/vit-age-classifier", |
|
"microsoft/resnet-50", |
|
"Falconsai/nsfw_image_detection", |
|
"cafeai/cafe_aesthetic", |
|
"microsoft/resnet-18", |
|
"microsoft/resnet-34", |
|
"microsoft/resnet-101", |
|
"microsoft/resnet-152", |
|
"microsoft/swin-tiny-patch4-window7-224", |
|
"-- Reinstated on testing--", |
|
"microsoft/beit-base-patch16-224-pt22k-ft22k", |
|
"-- 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() |
