add select feature

Author: danielsyahputra

App.py

@@ -16,6 +16,7 @@ def load_settings():
     st.markdown('## Auto Tagging for Fashion Retail Using Multi-label Image Classification')
 
     # Sidebar
+    global sidebar
     sidebar = st.sidebar
 
     with open("assets/test.jpg", "rb") as file:
@@ -35,13 +36,15 @@ def main():
     if uploaded_file is not None:
         image = Image.open(uploaded_file)
 
+        num_of_tags = sidebar.radio("Number of Tags", (3, 4, 5), index=2)
+
         col1, col2 = st.columns(2)
         with col1:
             st.markdown('<p style="text-align: center;">Your Image</p><hr>', unsafe_allow_html=True)
             st.image(image, caption="Fashion Image")
         with col2:
             st.markdown('<p style="text-align: center;">Image Tags</p><hr>', unsafe_allow_html=True)
-            tags = predictor.predict(image=image)
+            tags = predictor.predict(image=image, num_of_tags=num_of_tags)
             tags = st_tags(label="Tags:", text="Add more", value=tags)
 
 if __name__=="__main__":

utils.py

@@ -55,16 +55,16 @@ def load_model(self):
         model.load_state_dict(torch.load(f="./model/parameters.pth", map_location=torch.device('cpu')))
         self.model = model
 
-    def predict(self, image):
+    def predict(self, image, num_of_tags=5):
         image_transformed = self.transform(image)
         batch_image = torch.unsqueeze(image_transformed, 0)
         self.model.eval()
 
         # Inference
         output = self.model(batch_image).detach().numpy()
         preds = output[0]
-        top5 =  np.sort(preds)[::-1][min(4, len(preds)-1)]
-        preds[preds < top5] = 0
-        preds[preds >= top5] = 1
+        top_n_tags =  np.sort(preds)[::-1][min(num_of_tags - 1, len(preds)-1)]
+        preds[preds < top_n_tags] = 0
+        preds[preds >= top_n_tags] = 1
         tags = self.binarizer.inverse_transform(np.array([preds]))[0][:]
         return tags