Upload 15 files

tagger.py

@@ -59,26 +59,7 @@ def load_dict_from_csv(filename):
     return dict
 
 
-def get_series_dict():
-    import re
-    with open('characterfull.txt', 'r') as f:
-        lines = f.readlines()
-    series_dict = {}
-    for line in lines:
-        parts = line.strip().split(', ')
-        if len(parts) >= 3:
-            name = parts[-2].replace("\\", "")
-        if name.endswith(")"):
-            names = name.split("(")
-            character_name = "(".join(names[:-1])
-            if character_name.endswith(" "):
-                name = character_name[:-1]
-        series = re.sub(r'\\[()]', '', parts[-1])
-        series_dict[name] = series
-    return series_dict
-
-
-anime_series_dict = get_series_dict()
+anime_series_dict = load_dict_from_csv('character_series_dict.csv')
 
 
 def character_list_to_series_list(character_list):
@@ -248,7 +229,7 @@ def insert_recom_prompt(prompt: str = "", neg_prompt: str = "", type: str = "Non
     def list_uniq(l):
         return sorted(set(l), key=l.index)
 
-    animagine_ps = to_list("masterpiece, best quality, very aesthetic, absurdres")
+    animagine_ps = to_list("anime artwork, anime style, key visual, vibrant, studio anime, highly detailed, masterpiece, best quality, very aesthetic, absurdres")
     animagine_nps = to_list("lowres, (bad), text, error, fewer, extra, missing, worst quality, jpeg artifacts, low quality, watermark, unfinished, displeasing, oldest, early, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract]")
     pony_ps = to_list("source_anime, score_9, score_8_up, score_7_up, masterpiece, best quality, very aesthetic, absurdres")
     pony_nps = to_list("source_pony, source_furry, source_cartoon, score_6, score_5, score_4, busty, ugly face, mutated hands, low res, blurry face, black and white, the simpsons, overwatch, apex legends")
@@ -289,11 +270,11 @@ def remove_specific_prompt(input_prompt: str = "", keep_tags: str = "all"):
         return p.search(tag)
 
     un_tags = ['solo']
-    group_list = ['metatags', 'other-lists', 'artists', 'more-2', 'objects', 'other-groups', 'games', 'plants', 'real-world', 'more-1', 'genres', 'characters', 'image-style', 'see-also', 'creatures', 'sex', 'attire-accessories', 'body']
+    group_list = ['groups', 'body_parts', 'attire', 'posture', 'objects', 'creatures', 'locations', 'disambiguation_pages', 'commonly_misused_tags', 'phrases', 'verbs_and_gerunds', 'subjective', 'nudity', 'sex_objects', 'sex', 'sex_acts', 'image_composition', 'artistic_license', 'text', 'year_tags', 'metatags']
     keep_group_dict = {
-        "body": ['characters', 'body'],
-        "dress": ['characters', 'attire-accessories', 'body'],
-        "all": ['metatags', 'other-lists', 'artists', 'more-2', 'objects', 'other-groups', 'games', 'plants', 'real-world', 'more-1', 'genres', 'characters', 'image-style', 'see-also', 'creatures', 'sex', 'attire-accessories', 'body']
+        "body": ['groups', 'body_parts'],
+        "dress": ['groups', 'body_parts', 'attire'],
+        "all": group_list,
     }
 
     def is_necessary(tag, keep_tags, group_dict):
@@ -309,7 +290,7 @@ def remove_specific_prompt(input_prompt: str = "", keep_tags: str = "all"):
             return True
     
     if keep_tags == "all": return input_prompt
-    keep_group = keep_group_dict.get(keep_tags, ['characters', 'body'])
+    keep_group = keep_group_dict.get(keep_tags, keep_group_dict["body"])
     explicit_group = list(set(group_list) ^ set(keep_group))
 
     tags = input_prompt.split(",") if input_prompt else []
@@ -329,6 +310,59 @@ def remove_specific_prompt(input_prompt: str = "", keep_tags: str = "all"):
     return output_prompt
 
 
+def sort_taglist(tags: list[str]):
+    if not tags: return []
+    character_tags: list[str] = []
+    series_tags: list[str] = []
+    people_tags: list[str] = []
+    group_list = ['groups', 'body_parts', 'attire', 'posture', 'objects', 'creatures', 'locations', 'disambiguation_pages', 'commonly_misused_tags', 'phrases', 'verbs_and_gerunds', 'subjective', 'nudity', 'sex_objects', 'sex', 'sex_acts', 'image_composition', 'artistic_license', 'text', 'year_tags', 'metatags']
+    group_tags = {}
+    other_tags: list[str] = []
+    rating_tags: list[str] = []
+
+    group_dict = tag_group_dict
+    group_set = set(group_dict.keys())
+    character_set = set(anime_series_dict.keys())
+    series_set = set(anime_series_dict.values())
+    rating_set = set(DANBOORU_TO_E621_RATING_MAP.keys()) | set(DANBOORU_TO_E621_RATING_MAP.values())
+
+    for tag in tags:
+        tag = tag.strip().replace("_", " ")
+        if tag in PEOPLE_TAGS:
+            people_tags.append(tag)
+        elif tag in rating_set:
+            rating_tags.append(tag)
+        elif tag in group_set:
+            elem = group_dict[tag]
+            group_tags[elem] = group_tags[elem] + [tag] if elem in group_tags else [tag]
+        elif tag in character_set:
+            character_tags.append(tag)
+        elif tag in series_set:
+            series_tags.append(tag)
+        else:
+            other_tags.append(tag)
+
+    output_group_tags: list[str] = []
+    for k in group_list:
+        output_group_tags.extend(group_tags.get(k, []))
+
+    rating_tags = [rating_tags[0]] if rating_tags else []
+    rating_tags = ["explicit, nsfw"] if rating_tags and rating_tags[0] == "explicit" else rating_tags
+
+    output_tags = character_tags + series_tags + people_tags + output_group_tags + other_tags + rating_tags
+    
+    return output_tags
+
+
+def sort_tags(tags: str):
+    if not tags: return ""
+    taglist: list[str] = []
+    for tag in tags.split(","):
+        taglist.append(tag.strip())
+    taglist = list(filter(lambda x: x != "", taglist))
+    return ", ".join(sort_taglist(taglist))
+
+
 def postprocess_results(results: dict[str, float], general_threshold: float, character_threshold: float):
     results = {
         k: v for k, v in sorted(results.items(), key=lambda item: item[1], reverse=True)
@@ -400,7 +434,13 @@ def predict_tags(image: Image.Image, general_threshold: float = 0.3, character_t
 
     return output_series_tag, ", ".join(character.keys()), prompt, gr.update(interactive=True),
 
-            
+
+def predict_tags_wd(image: Image.Image, input_tags: str, algo: list[str], general_threshold: float = 0.3, character_threshold: float = 0.8):
+    if algo and not "Use WD Tagger" in algo:
+        return "", "", input_tags, gr.update(interactive=True),
+    return predict_tags(image, general_threshold, character_threshold)
+
+
 def compose_prompt_to_copy(character: str, series: str, general: str):
     characters = character.split(",") if character else []
     serieses = series.split(",") if series else []