modify od fashion

pages/object_detection.py

@@ -4,14 +4,24 @@ import streamlit as st
 import matplotlib.pyplot as plt
 import pandas as pd
 import numpy as np
-import altair as alt
+#import altair as alt
+import plotly.express as px
+
 
 from PIL import Image
 from transformers import YolosFeatureExtractor, YolosForObjectDetection
 from torchvision.transforms import ToTensor, ToPILImage
+#from utils import load_model_huggingface
+
 
 st.set_page_config(layout="wide")
 
+@st.cache_data(ttl=3600, show_spinner=False)
+def load_model(feature_extractor_url, model_url):
+    feature_extractor_ = YolosFeatureExtractor.from_pretrained(feature_extractor_url)
+    model_ = YolosForObjectDetection.from_pretrained(model_url)
+    return feature_extractor_, model_
+    
 
 def rgb_to_hex(rgb):
     """Converts an RGB tuple to an HTML-style Hex string."""
@@ -76,7 +86,7 @@ def plot_results(pil_img, prob, boxes):
     
     plt.savefig("results_od.png", 
             bbox_inches ="tight") 
-    #plt.show()
+    plt.show()
     st.image("results_od.png")
 
     return colors_used
@@ -112,15 +122,23 @@ def visualize_probas(probas, threshold, colors):
     top_label_df["colors"] = colors
     top_label_df.sort_values(by=["proba"], ascending=False, inplace=True)
 
-    st.dataframe(top_label_df.drop(columns=["colors"]))
+    #st.dataframe(top_label_df.drop(columns=["colors"]))
     
     mode_func = lambda x: x.mode().iloc[0]
     top_label_df_agg = top_label_df.groupby("label").agg({"proba":"mean", "colors":mode_func})
     top_label_df_agg = top_label_df_agg.reset_index().sort_values(by=["proba"], ascending=False)
+    top_label_df_agg.columns = ["Item","Score","Colors"]
+
+    color_map = dict(zip(top_label_df_agg["Item"].to_list(),
+                         top_label_df_agg["Colors"].to_list()))
+
+    fig = px.bar(top_label_df_agg, y='Item', x='Score',
+                 color="Item", title="Probability scores")
+    st.plotly_chart(fig, use_container_width=True)
 
-    chart = alt.Chart(top_label_df_agg).mark_bar().encode(x="proba", y="label", 
-                                                          color=alt.Color('colors:N', scale=None)).interactive()
-    #st.altair_chart(chart)
+    # chart = alt.Chart(top_label_df_agg).mark_bar().encode(x="proba", y="label", 
+    #                                                       color=alt.Color('colors:N', scale=None)).interactive()
+    # st.altair_chart(chart)
 
 
 
@@ -156,34 +174,38 @@ st.markdown("""Common applications of Object Detection include:
 st.markdown("  ")
 st.divider()
 
-st.markdown("### Fashion object detection 👗")
-st.markdown(""" The following example showcases the use of an **Object detection algorithm** for clothing items/features on fashion images. <br>
-            This use case can be seen as an application of AI models for Fashion and E-commerce. <br>
-            """, unsafe_allow_html=True)
+st.markdown("## Fashion Object Detection 👗")
+# st.info("""This use case showcases the application of **Object detection** to detect clothing items/features on images. <br>
+#             The images used were gathered from Dior's""")
+st.info("""In this use case, we are going to identify and locate different articles of clothings, as well as finer details such as a collar or pocket using an object detection AI model. 
+        The images used were taken from **Dior's 2020 Fall Women Fashion Show**.""")
+
+st.markdown("  ")
 
-st.image("images/od_fashion.jpg", width=700)
+images_dior = [os.path.join("data/dior_show",url) for url in os.listdir("data/dior_show") if url != "results"]
+columns_img = st.columns(4)
+for img, col in zip(images_dior,columns_img):
+    with col:
+        st.image(img)
 
-#images_dior = [os.path.join("data/dior_show",url) for url in os.listdir("data/dior_show") if url != "results"]
-#st.image(images_dior, width=250, caption=[file for file in os.listdir("data/dior_show") if file != "results"])
 
 st.markdown("  ")
-#st.markdown("##### Select an image")
 
 
 ############## SELECT AN IMAGE ###############
 
-st.markdown("#### Step 1: Select an image")
-st.info("""First, select the image that you wish to use the Object detecion model on.""") 
-st.markdown("**Note:** The model was trained to detect clothing items on a single person. If your image has more than individuals, the model will ignore one of them in its detection.")
+st.markdown("#### Select an image 🖼️")
+#st.markdown("""**Select an image that you wish to run the Object Detection model on.**""") 
+
 
 image_ = None
 select_image_box = st.radio(
-    "",
+    "**Select the image you wish to run the model on**",
     ["Choose an existing image", "Load your own image"],
-    index=None, label_visibility="collapsed")
+    index=None,)# #label_visibility="collapsed")
 
 if select_image_box == "Choose an existing image":
-    fashion_images_path = r"data/pinterest"
+    fashion_images_path = r"data/dior_show"
     list_images = os.listdir(fashion_images_path)
     image_ = st.selectbox("", list_images, label_visibility="collapsed")
     
@@ -198,6 +220,8 @@ elif select_image_box == "Load your own image":
     
     st.warning("""**Note**: The model tends to perform better with images of people/clothing items facing forward. 
            Choose this type of image if you want optimal results.""")
+    st.warning("""**Note:** The model was trained to detect clothing items on a single person. 
+               If your image contains more than one person, the model won't detect the items of the other persons.""")
 
     if image_ is not None:
         st.image(Image.open(image_), width=300)
@@ -216,7 +240,7 @@ cats = ['shirt, blouse', 'top, t-shirt, sweatshirt', 'sweater', 'cardigan', 'jac
 
 dict_cats = dict(zip(np.arange(len(cats)), cats))
 
-st.markdown("#### Step 2: Choose the elements you want to detect")
+st.markdown("#### Choose the elements you want to detect 👉")
 
 # Select one or more elements to detect
 container = st.container()
@@ -239,21 +263,31 @@ st.markdown("  ")
 
 ############## SELECT A THRESHOLD ###############
 
-st.markdown("#### Step 3: Select a threshold")
+st.markdown("#### Define a threshold for predictions 🔎")
+st.markdown("""Object detection models assign to each element detected a **probability score**. <br>
+            This score represents the model's belief in the accuracy of its prediction for a specific object.
+            """, unsafe_allow_html=True)
 
-st.markdown("""Finally, select a threshold for the model. 
-            The threshold helps you decide how confident you want your model to be with its predictions. 
-            Elements that were identified with a lower probability than the given threshold will be ignored in the final results.""")
+st.warning("**Note:** Objects that are assigned a lower score than the chosen threshold will be ignored in the final results.")
 
-threshold = st.slider('**Select a threshold**', min_value=0.0, step=0.05, max_value=1.0, value=0.75, label_visibility="collapsed")
-# min_value=0.000000, step=0.000001, max_value=0.0005, value=0.0000045, format="%f"
 
-if threshold < 0.6:
-    st.warning("""**Warning**: Selecting a low threshold (below 0.6) could lead the model to make errors and detect too many objects.""")
+_, col, _ = st.columns([0.2,0.6,0.2])
+with col:
+    st.image("images/probability_od.png", caption="Example of object detection with probability scores")
 
-st.write("You've selected a threshold at", threshold)
+st.markdown(" ")
+
+st.markdown("**Select a threshold** ")
+
+# st.warning("""**Note**: The threshold helps you decide how confident you want your model to be with its predictions. 
+#             Elements that are identified with a lower probability than the given threshold will be ignored in the final results.""")
 
+threshold = st.slider('**Select a threshold**', min_value=0.5, step=0.05, max_value=1.0, value=0.75, label_visibility="collapsed")
 
+if threshold < 0.6:
+    st.error("""**Warning**: Selecting a low threshold (below 0.6) could lead the model to make errors and detect too many objects.""")
+
+st.write("You've selected a threshold at", threshold)
 st.markdown("  ")
 
 
@@ -269,29 +303,37 @@ if run_model:
             image = fix_channels(ToTensor()(image))
 
             ## LOAD OBJECT DETECTION MODEL
-            MODEL_NAME = "valentinafeve/yolos-fashionpedia"
-            feature_extractor = YolosFeatureExtractor.from_pretrained('hustvl/yolos-small')
-            model = YolosForObjectDetection.from_pretrained(MODEL_NAME)
+            FEATURE_EXTRACTOR_PATH = "hustvl/yolos-small"
+            MODEL_PATH = "valentinafeve/yolos-fashionpedia"
+            feature_extractor, model = load_model(FEATURE_EXTRACTOR_PATH, MODEL_PATH)
+            # feature_extractor = YolosFeatureExtractor.from_pretrained('hustvl/yolos-small')
+            # model = YolosForObjectDetection.from_pretrained(MODEL)
 
             # RUN MODEL ON IMAGE
             inputs = feature_extractor(images=image, return_tensors="pt")
             outputs = model(**inputs)
             probas, keep = return_probas(outputs, threshold)
 
+            st.markdown("#### See the results ☑️")
+
             # PLOT BOUNDING BOX AND BARS/PROBA
             col1, col2 = st.columns(2)
             with col1:
-                st.markdown("##### Bounding box results")
+                #st.markdown("**Bounding box results**")
                 bboxes_scaled = rescale_bboxes(outputs.pred_boxes[0, keep].cpu(), image.size)
                 colors_used = plot_results(image, probas[keep], bboxes_scaled)
             
             with col2: 
-                visualize_probas(probas, threshold, colors_used)
+                #st.markdown("**Probability scores**")
+                if not any(keep.tolist()):
+                    st.error("""No objects were detected on the image. 
+                             Decrease your threshold or choose differents items to detect.""")
+                else:
+                    visualize_probas(probas, threshold, colors_used)
             
-            st.info("Done")
 
     else:
-        st.warning("You must select an **image**, **elements to detect** and a **threshold** to run the model !")
+        st.error("You must select an **image**, **elements to detect** and a **threshold** to run the model !")