Upload app.py

app.py

@@ -14,6 +14,12 @@ DESCRIPTION = """
 implications_list_path = './implications_list.json'
 related_feature_path = './related_feature.json'
 
+with open(related_feature_path, 'r') as f:
+    related_feature_list = json.load(f)
+
+with open(implications_list_path, 'r') as f:
+    implications_list = json.load(f)
+
 #HF_TOKEN = os.environ["HF_TOKEN"]
 
 # Dataset v3 series of models:
@@ -21,6 +27,7 @@ SWINV2_MODEL_DSV3_REPO = "SmilingWolf/wd-swinv2-tagger-v3"
 CONV_MODEL_DSV3_REPO = "SmilingWolf/wd-convnext-tagger-v3"
 VIT_MODEL_DSV3_REPO = "SmilingWolf/wd-vit-tagger-v3"
 VIT_LARGE_MODEL_DSV3_REPO = "SmilingWolf/wd-vit-large-tagger-v3"
+EVA02_LARGE_MODEL_DSV3_REPO = "SmilingWolf/wd-eva02-large-tagger-v3"
 
 # Dataset v2 series of models:
 MOAT_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-moat-tagger-v2"
@@ -61,8 +68,7 @@ def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser()
     parser.add_argument("--score-slider-step", type=float, default=0.05)
     parser.add_argument("--score-general-threshold", type=float, default=0.4)
-    parser.add_argument("--score-character-threshold", type=float, default=0.9)
-    parser.add_argument("--character_string", type=str)
+    parser.add_argument("--score-character-threshold", type=float, default=0.8)
     parser.add_argument("--share", action="store_true")
     return parser.parse_args()
 
@@ -170,6 +176,7 @@ class Predictor:
         character_thresh,
         character_mcut_enabled,
         character_string,
+        character_output
     ):
         self.load_model(model_repo)
 
@@ -195,13 +202,6 @@ class Predictor:
         general_res = [x for x in general_names if x[1] > general_thresh]
         general_res = dict(general_res)
 
-        with open(related_feature_path, 'r') as f:
-            related_feature_list = json.load(f)
-
-
-        with open(implications_list_path, 'r') as f:
-            implications_list = json.load(f)
-
         to_delete = set()
         for key in general_res.keys():
             if key in implications_list:
@@ -221,16 +221,30 @@ class Predictor:
         character_res = [x for x in character_names if x[1] > character_thresh]
         character_res = dict(character_res)
 
+        character_strings = sorted(
+            character_res.items(),
+            key=lambda x: x[1],
+            reverse=True,
+        )
+        character_strings = [x[0] for x in character_strings]
+
         sorted_general_strings = sorted(
             general_res.items(),
             key=lambda x: x[1],
             reverse=True,
         )
 
-        character_list = character_string.lower().split(', ')
+        character_list = []
+        if character_string != '':
+            character_list = character_string.lower().split(', ')
+
+        if character_output:
+            character_combined = character_list + character_strings
+        else:
+            character_combined = character_list
 
         feature_delete_list = []
-        for tag in character_list:
+        for tag in character_combined:
             if tag in related_feature_list:  
                 feature_delete_list.extend(related_feature_list[tag])
 
@@ -238,6 +252,8 @@ class Predictor:
 
         sorted_general_strings = [x for x in sorted_general_strings if x not in feature_delete_list]
 
+        sorted_general_strings = character_combined + sorted_general_strings
+
         sorted_general_strings = [x.replace("_", " ") if x not in kaomojis else x for x in sorted_general_strings]
 
         sorted_general_strings = (
@@ -246,7 +262,6 @@ class Predictor:
 
         return sorted_general_strings, rating, character_res, general_res
 
-
 def main():
     args = parse_args()
 
@@ -257,6 +272,7 @@ def main():
         CONV_MODEL_DSV3_REPO,
         VIT_MODEL_DSV3_REPO,
         VIT_LARGE_MODEL_DSV3_REPO,
+        EVA02_LARGE_MODEL_DSV3_REPO,
     ]
 
     with gr.Blocks(title=TITLE) as demo:
@@ -306,6 +322,11 @@ def main():
                         label= "Character",
                         scale=3,
                         )
+                        character_output = gr.Checkbox(
+                            value=True,
+                            label="Use Output (characters)",
+                            scale=1,
+                        )
                     with gr.Row():
                         clear = gr.ClearButton(
                             components=[
@@ -340,6 +361,7 @@ def main():
                 character_thresh,
                 character_mcut_enabled,
                 character_string,
+                character_output
             ],
             outputs=[sorted_general_strings, rating, character_res, general_res],
         )