feat: run all models and display results in table instead of dropdown

app.py

@@ -1,164 +1,160 @@
 import streamlit as st
 from transformers import pipeline
 from PIL import Image
+import pandas as pd
 
-MODEL_1 = "google/vit-base-patch16-224"
 MIN_ACEPTABLE_SCORE = 0.1
 MAX_N_LABELS = 5
-MODEL_2 = "nateraw/vit-age-classifier"
+
 MODELS = [
-            "-- General Image Classification --",
-            "google/vit-base-patch16-224",
-            "microsoft/resnet-50",
-            "microsoft/resnet-18",
-            "microsoft/resnet-34",
-            "microsoft/resnet-101",
-            "microsoft/resnet-152",
-            "microsoft/swin-tiny-patch4-window7-224",
-            "microsoft/swinv2-base-patch4-window16-256",
-            "microsoft/beit-base-patch16-224-pt22k-ft22k",
-            "facebook/convnext-large-224",
-            "facebook/convnext-base-224-22k-1k",
-            "facebook/convnext-tiny-224",
-            "nvidia/mit-b0",
-            "timm/resnet50.a1_in1k",
-            "timm/tf_efficientnetv2_s.in21k",
-            "timm/convnext_tiny.fb_in22k",
-            "vit-base-patch16-224-in21k",
-            "facebook/deit-base-distilled-patch16-224 << new >>",
-            "WinKawaks/vit-tiny-patch16-224 << new >>",
-
-            "-- Age Classification --",
-            "nateraw/vit-age-classifier",
-
-            "-- NSFW Detection --",
-            "Falconsai/nsfw_image_detection",
-            "LukeJacob2023/nsfw-image-detector",
-            "carbon225/vit-base-patch16-224-hentai",
-            "Marqo/nsfw-image-detection-384 << new >>",
-
-            "-- Aesthetic/Art Classification --",
-            "cafeai/cafe_aesthetic",
-            "shadowlilac/aesthetic-shadow",
-            "pixai-labs/pixai-tagger-v0.9 << new >>",
-
-            "-- Face/Emotion Classification --",
-            "trpakov/vit-face-expression",
-            "RickyIG/emotion_face_image_classification",
-            "rizvandwiki/gender-classification",
-
-            "-- Food Classification --",
-            "nateraw/food",
-            "BinhQuocNguyen/food-recognition-model << new >>",
-
-            "-- Medical/Dermatology --",
-            "google/derm-foundation << new >>",
-            "google/cxr-foundation << new >>",
-            "Anwarkh1/Skin_Cancer-Image_Classification << new >>",
-
-            "-- AI vs Human Detection --",
-            "Ateeqq/ai-vs-human-image-detector << new >>",
-            "umm-maybe/AI-image-detector << new >>",
-
-            "-- Deepfake Detection --",
-            "not-lain/deepfake",
-
-            "-- Anime/Manga Classification --",
-            #"Readidno/anime.mili << new >>",  # Not working - missing model_type
-
-            "-- Human Activity Recognition --",
-            "DunnBC22/vit-base-patch16-224-in21k_Human_Activity_Recognition",
-
-            "-- Clothing/Fashion --",
-            "aalonso-developer/vit-base-patch16-224-in21k-clothing-classifier",
-
-            "-- Real Estate --",
-            "andupets/real-estate-image-classification",
-
-            "-- Satellite/Remote Sensing --",
-            "FatihC/swin-tiny-patch4-window7-224-finetuned-eurosat-watermark",
-
-            "-- Car Classification --",
-            "lamnt2008/car_brands_classification << new >>",
-
-            "-- Document Classification --",
-            "docling-project/DocumentFigureClassifier-v2.5 << new >>",
-
-            "-- EfficientNet (timm) --",
-            "timm/efficientnet_b0.ra_in1k << new >>",
-            "timm/mobilenetv3_large_100.ra_in1k",
-            "timm/mobilenetv3_small_100.lamb_in1k << new >>",
-
-            "-- Experimental/Future --",
-            "#q-future/one-align",
-        ]
+            ("google/vit-base-patch16-224", "General Image Classification"),
+            ("microsoft/resnet-50", "General Image Classification"),
+            ("microsoft/resnet-18", "General Image Classification"),
+            ("microsoft/resnet-34", "General Image Classification"),
+            ("microsoft/resnet-101", "General Image Classification"),
+            ("microsoft/resnet-152", "General Image Classification"),
+            ("microsoft/swin-tiny-patch4-window7-224", "General Image Classification"),
+            ("microsoft/swinv2-base-patch4-window16-256", "General Image Classification"),
+            ("microsoft/beit-base-patch16-224-pt22k-ft22k", "General Image Classification"),
+            ("facebook/convnext-large-224", "General Image Classification"),
+            ("facebook/convnext-base-224-22k-1k", "General Image Classification"),
+            ("facebook/convnext-tiny-224", "General Image Classification"),
+            ("nvidia/mit-b0", "General Image Classification"),
+            ("timm/resnet50.a1_in1k", "General Image Classification"),
+            ("timm/tf_efficientnetv2_s.in21k", "General Image Classification"),
+            ("timm/convnext_tiny.fb_in22k", "General Image Classification"),
+            ("google/vit-base-patch16-224-in21k", "General Image Classification"),
+            ("facebook/deit-base-distilled-patch16-224", "General Image Classification"),
+            ("WinKawaks/vit-tiny-patch16-224", "General Image Classification"),
 
-def classify(image, model):
-    model_name = model.replace(" << new >>", "")
-    classifier = pipeline("image-classification", model=model_name)
-    result= classifier(image)
-    return result
+            ("nateraw/vit-age-classifier", "Age Classification"),
 
-def save_result(result):
-    st.write("In the future, this function will save the result in a database.")
+            ("Falconsai/nsfw_image_detection", "NSFW Detection"),
+            ("LukeJacob2023/nsfw-image-detector", "NSFW Detection"),
+            ("carbon225/vit-base-patch16-224-hentai", "NSFW Detection"),
+            ("Marqo/nsfw-image-detection-384", "NSFW Detection"),
 
-def print_result(result):
+            ("cafeai/cafe_aesthetic", "Aesthetic/Art Classification"),
+            ("shadowlilac/aesthetic-shadow", "Aesthetic/Art Classification"),
+            ("pixai-labs/pixai-tagger-v0.9", "Aesthetic/Art Classification"),
 
-    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'])
+            ("trpakov/vit-face-expression", "Face/Emotion Classification"),
+            ("RickyIG/emotion_face_image_classification", "Face/Emotion Classification"),
+            ("rizvandwiki/gender-classification", "Face/Emotion Classification"),
 
-    st.write(f"comulative_discarded_score:")
-    st.progress(comulative_discarded_score)
-    st.write(comulative_discarded_score)
-    
+            ("nateraw/food", "Food Classification"),
+            ("BinhQuocNguyen/food-recognition-model", "Food Classification"),
 
+            ("google/derm-foundation", "Medical/Dermatology"),
+            ("google/cxr-foundation", "Medical/Dermatology"),
+            ("Anwarkh1/Skin_Cancer-Image_Classification", "Medical/Dermatology"),
 
-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.markdown(":white_check_mark: **:green[22 new models added!]** - Including Medical, AI vs Human detection, Anime classification and more.")
-    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("If you like this project, please consider liking it or buying me a coffee. It will help me to keep working on this and other projects. Thank you!")
+            ("Ateeqq/ai-vs-human-image-detector", "AI vs Human Detection"),
+            ("umm-maybe/AI-image-detector", "AI vs Human Detection"),
 
-    # Buy me a Coffee Setup
-    bmc_link = "https://www.buymeacoffee.com/nuno.tome" 
-    # image_url = "https://helloimjessa.files.wordpress.com/2021/06/bmc-button.png?w=150" # Image URL
-    image_url = "https://i.giphy.com/RETzc1mj7HpZPuNf3e.webp" # Image URL
-        
-    image_size = "150px" # Image size
-    #image_link_markdown = f"<img src='{image_url}' width='25%'>"
-    image_link_markdown = f"[![Buy Me a Coffee]({image_url})]({bmc_link})"
+            ("not-lain/deepfake", "Deepfake Detection"),
+
+            ("DunnBC22/vit-base-patch16-224-in21k_Human_Activity_Recognition", "Human Activity Recognition"),
+
+            ("aalonso-developer/vit-base-patch16-224-in21k-clothing-classifier", "Clothing/Fashion"),
+
+            ("andupets/real-estate-image-classification", "Real Estate"),
 
-    #image_link_markdown = f"[![Buy Me a Coffee]({image_url})]({bmc_link})" # Create a clickable image link
+            ("FatihC/swin-tiny-patch4-window7-224-finetuned-eurosat-watermark", "Satellite/Remote Sensing"),
+
+            ("lamnt2008/car_brands_classification", "Car Classification"),
+
+            ("docling-project/DocumentFigureClassifier-v2.5", "Document Classification"),
+
+            ("timm/efficientnet_b0.ra_in1k", "EfficientNet"),
+            ("timm/mobilenetv3_large_100.ra_in1k", "EfficientNet"),
+            ("timm/mobilenetv3_small_100.lamb_in1k", "EfficientNet"),
+        ]
 
-    st.markdown(image_link_markdown, unsafe_allow_html=True) # Display the image link
-    # Buy me a Coffee Setup
+def classify(image, model_name):
+    classifier = pipeline("image-classification", model=model_name)
+    result = classifier(image)
+    return result
+
+def format_results(results):
+    labels = []
+    scores = []
+    for r in results[:MAX_N_LABELS]:
+        if r['score'] >= MIN_ACEPTABLE_SCORE:
+            labels.append(r['label'])
+            scores.append(f"{r['score']:.2%}")
+    return "<br>".join(labels), "<br>".join(scores)
+
+def main():
+    st.title("Image Classification - Compare All Models")
+    st.write("This app runs ALL image classification models and displays results in a table.")
+    st.markdown(":white_check_mark: **:green[Run all models at once!]**")
     
-    #st.markdown("<img src='https://helloimjessa.files.wordpress.com/2021/06/bmc-button.png?w=1024' width='15%'>", unsafe_allow_html=True)
+    bmc_link = "https://www.buymeacoffee.com/nuno.tome" 
+    image_url = "https://i.giphy.com/RETzc1mj7HpZPuNf3e.webp" 
+    image_link_markdown = f"[![Buy Me a Coffee]({image_url})]({bmc_link})"
+    st.markdown(image_link_markdown, unsafe_allow_html=True)
   
     input_image = st.file_uploader("Upload Image")
-    shosen_model = st.selectbox("Select the model to use",  MODELS)
-    
-    
+
     if input_image is not None:
         image_to_classify = Image.open(input_image)
-        st.image(image_to_classify, caption="Uploaded Image")
-        if st.button("Classify"):
-            image_to_classify = Image.open(input_image)
-            classification_obj1 =[]
-            #avable_models = st.selectbox
+        st.image(image_to_classify, caption="Uploaded Image", use_container_width=True)
+        
+        if st.button("Run All Models", type="primary"):
+            results_data = []
+            progress_bar = st.progress(0)
+            status_text = st.empty()
             
-            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()
+            for i, (model_name, category) in enumerate(MODELS):
+                status_text.text(f"Running model {i+1}/{len(MODELS)}: {model_name}")
+                try:
+                    classification_result = classify(image_to_classify, model_name)
+                    labels, scores = format_results(classification_result)
+                    results_data.append({
+                        "Model": model_name,
+                        "Category": category,
+                        "Top Labels": labels,
+                        "Scores": scores
+                    })
+                except Exception as e:
+                    results_data.append({
+                        "Model": model_name,
+                        "Category": category,
+                        "Top Labels": f"Error: {str(e)[:50]}",
+                        "Scores": "-"
+                    })
+                progress_bar.progress((i + 1) / len(MODELS))
+            
+            status_text.text("Done!")
+            
+            if results_data:
+                df = pd.DataFrame(results_data)
+                st.subheader(f"Results ({len(results_data)} models)")
+                
+                st.markdown("""
+                <style>
+                .dataframe {font-size: 12px;}
+                </style>
+                """, unsafe_allow_html=True)
+                
+                st.dataframe(
+                    df,
+                    use_container_width=True,
+                    hide_index=True,
+                    column_config={
+                        "Model": st.column_config.TextColumn("Model", width="medium"),
+                        "Category": st.column_config.TextColumn("Category", width="small"),
+                        "Top Labels": st.column_config.TextColumn("Top Labels", width="large"),
+                        "Scores": st.column_config.TextColumn("Scores", width="medium"),
+                    }
+                )
+                
+                csv = df.to_csv(index=False).encode('utf-8')
+                st.download_button(
+                    "Download Results CSV",
+                    csv,
+                    "classification_results.csv",
+                    "text/csv",
+                    key='download-csv'
+                )