Add underline button

app.py

@@ -1,6 +1,4 @@
 import argparse
-import os
-
 import gradio as gr
 import huggingface_hub
 import numpy as np
@@ -158,7 +156,8 @@ class Predictor:
 
         return np.expand_dims(image_array, axis=0)
     
-    def tag_dict_to_sorted_string(self, dict_res: dict, sort_by_confidence, descending, escape_parens):
+    def tag_dict_to_sorted_string(self, dict_res: dict, sort_by_confidence, descending,
+                                  remove_underlines, escape_parens):
         """Custom function: Sort tag dict by confidence/alphabetically"""
         if sort_by_confidence:
             _sorted_list = sorted(
@@ -171,9 +170,12 @@ class Predictor:
                 dict_res.items(),
                 reverse=descending
             )
-        _sorted_string = ", ".join([x[0] for x in _sorted_list])
+        if remove_underlines:
+            _sorted_string = ", ".join([x[0] for x in _sorted_list])
+        else:  # Add back underlines
+            _sorted_string = ", ".join([x[0].replace(" ", "_") for x in _sorted_list])
         if escape_parens:
-            _sorted_string = _sorted_string.replace("(", "\(").replace(")", "\)")
+            _sorted_string = _sorted_string.replace("(", "\\(").replace(")", "\\)")
         return _sorted_string
 
     def predict(
@@ -186,6 +188,7 @@ class Predictor:
         character_mcut_enabled,
         sort_by_confidence_enabled,
         sort_descending_enabled,
+        remove_underline_enabled,
         escape_parens_enabled
     ):
         self.load_model(model_repo)
@@ -227,12 +230,14 @@ class Predictor:
             general_res,
             sort_by_confidence=sort_by_confidence_enabled,
             descending=sort_descending_enabled,
+            remove_underlines=remove_underline_enabled,
             escape_parens=escape_parens_enabled
         )
         sorted_character_strings = self.tag_dict_to_sorted_string(
             character_res,
             sort_by_confidence=sort_by_confidence_enabled,
             descending=sort_descending_enabled,
+            remove_underlines=remove_underline_enabled,
             escape_parens=escape_parens_enabled
         )
 
@@ -323,6 +328,11 @@ def main():
                             value=False,
                             label="Descending"
                         )
+                    with gr.Row():
+                        remove_underline_enabled = gr.Checkbox(
+                            value=True,
+                            label="Remove Tag Underlines"
+                        )
                         escape_parens_enabled = gr.Checkbox(
                             value=False,
                             label="Escape Parens"
@@ -332,7 +342,7 @@ def main():
                         show_copy_button=True
                     )
                     sorted_character_strings = gr.Textbox(
-                        label="Character (string)",
+                        label="Characters (string)",
                         show_copy_button=True
                     )
                     rating = gr.Label(label="Rating")
@@ -358,6 +368,7 @@ def main():
                 character_mcut_enabled,
                 sort_by_confidence_enabled,
                 sort_descending_enabled,
+                remove_underline_enabled,
                 escape_parens_enabled
             ],
             outputs=[sorted_general_strings, sorted_character_strings,
@@ -365,7 +376,7 @@ def main():
         )
 
     demo.queue(max_size=10)
-    demo.launch()
+    demo.launch(share=args.share)
 
 
 if __name__ == "__main__":